import dataclasses
from functools import cached_property

import torch


__all__ = []


def _make_empty(shape, dtype, device, when=True):
    # when=False returns a rank-matching empty (shape (0,)*len(shape)).
    return torch.empty(
        shape if when else (0,) * len(shape),
        device=device,
        dtype=dtype,
        requires_grad=False,
    )


def _make_zeros(shape, dtype, device, when=True):
    return torch.zeros(
        shape if when else (0,) * len(shape),
        device=device,
        dtype=dtype,
        requires_grad=False,
    )


def _linear_cross_entropy_batch_chunked_setup_context(ctx, inputs, output):
    (
        *_,
        allow_retain_graph,
        compute_input_grad,
        compute_linear_weight_grad,
        compute_linear_bias_grad,
    ) = inputs
    ctx.allow_retain_graph = allow_retain_graph
    ctx.compute_input_grad = compute_input_grad
    ctx.compute_linear_weight_grad = compute_linear_weight_grad
    ctx.compute_linear_bias_grad = compute_linear_bias_grad
    _, grad_input, grad_linear_weight, grad_linear_bias = output
    # ``.detach()`` the stashed grads: they are forward OUTPUTS, so they
    # carry ``grad_fn`` back to this custom-op node. Holding them as raw
    # ctx attributes would form a reference cycle (node -> ctx -> grad ->
    # grad_fn -> node) that refcounting can't break, leaking the tensors
    # until a GC pass -- which the CUDA mem-leak check flags. Detaching
    # drops the grad_fn link so the tensors free promptly when the loss's
    # graph is. Backward only multiplies them by the upstream grad, so the
    # lost autograd history is irrelevant (double-backward is unsupported).
    ctx._gi = grad_input.detach() if compute_input_grad else None
    ctx._gw = grad_linear_weight.detach() if compute_linear_weight_grad else None
    ctx._gb = grad_linear_bias.detach() if compute_linear_bias_grad else None


@dataclasses.dataclass
class _ChunkViews:
    """Per-iteration tensor views; each property picks the right operand
    or scratch from ctx dispatch flags so call sites read raw.
    """

    ctx: "_ChunkContext"
    bchunk_start: int
    bchunk_size: int
    input_chunk: torch.Tensor
    target_chunk: torch.Tensor
    weight_chunk: torch.Tensor
    logits: torch.Tensor
    input_chunk_acc: torch.Tensor

    @property
    def input(self) -> torch.Tensor:
        return (
            self.input_chunk_acc
            if self.ctx.forward_uses_acc_input
            else self.input_chunk
        )

    @property
    def linear_weight(self) -> torch.Tensor:
        ctx = self.ctx
        return (
            ctx.linear_weight
            if ctx.forward_uses_cuda_out_dtype
            else ctx.linear_weight_cast
        )

    @property
    def input_grad_logits(self) -> torch.Tensor:
        if self.ctx.input_grad_uses_logits_lw:
            return self.logits_downcast
        return self.logits

    @property
    def input_grad_linear_weight(self) -> torch.Tensor:
        ctx = self.ctx
        return (
            ctx.linear_weight
            if ctx.input_grad_uses_logits_lw
            else ctx.linear_weight_cast
        )

    @property
    def weight_grad_input(self) -> torch.Tensor:
        ctx = self.ctx
        if ctx.is_cuda and not ctx.weight_grad_mm_same_dtype:
            return self.input_chunk
        return self.input_chunk_acc

    @property
    def logits_upcast(self) -> torch.Tensor:
        # Non-CUDA mixed-dtype: copy into logits_acc_buf for matching-dtype mm.
        ctx = self.ctx
        if ctx.is_cuda or ctx.weight_grad_mm_same_dtype:
            return self.logits
        return ctx.logits_acc_buf.narrow(0, 0, self.bchunk_size).copy_(self.logits)

    @cached_property
    def grad_input_chunk(self) -> torch.Tensor:
        # Buf-and-copy commit happens in ctx.chunks() post-yield.
        ctx = self.ctx
        if ctx.alloc_input_grad_acc_buf:
            return ctx.input_grad_acc_buf.narrow(0, 0, self.bchunk_size)
        return ctx.grad_input.narrow(0, self.bchunk_start, self.bchunk_size)

    @cached_property
    def bias_grad_acc(self) -> torch.Tensor:
        # ``ctx.linear_bias_grad_chunk`` is an acc_dtype scratch that
        # accumulates the per-chunk sum + index_add at acc precision,
        # then commits to ``grad_linear_bias`` (input dtype) via the
        # ``ctx.chunks()`` post-yield hook -- mirrors
        # ``weight_grad_chunk``'s bulk-+-correction pattern for
        # ``grad_linear_weight``. On the non-acc path, accumulation
        # goes directly into ``grad_linear_bias`` (no commit, no
        # post-yield action).
        ctx = self.ctx
        if ctx.alloc_linear_bias_grad_chunk:
            return ctx.linear_bias_grad_chunk
        return ctx.grad_linear_bias

    @cached_property
    def logits_downcast(self) -> torch.Tensor:
        # Lazy: first access must come AFTER in-place mods to self.logits.
        ctx = self.ctx
        if ctx.loop_caches_logits_downcast:
            return self.logits.to(ctx.linear_weight.dtype)
        return self.logits


@dataclasses.dataclass
class _ChunkContext:
    """Per-call state for the chunked loop, built once via ``build``.
    Methods hide dtype/device/acc_policy dispatch behind single math ops;
    dispatch-free per-iter math is inlined into the loop body. Buffers
    dispatch decided are not needed are rank-matching empty tensors
    (``when=False`` in ``_make_*``) so the dataclass surface stays uniform.
    """

    dtype: torch.dtype
    num_batches: int
    in_features: int
    num_classes: int
    is_cuda: bool
    use_acc_dtype: bool

    acc_dtype: torch.dtype
    weight_chunk_dtype: torch.dtype
    grad_input_dtype: torch.dtype
    linear_weight_cast_dtype: torch.dtype

    input: torch.Tensor
    target: torch.Tensor  # uncorrected; per-iter slice is from corrected_target
    weight: torch.Tensor | None
    ignore_index: int
    reduction: str
    linear_weight: torch.Tensor
    linear_bias: torch.Tensor | None
    # reduction='none' backward only: per-sample upstream grad (N,). When
    # set, neg_weight_target folds it in so the shared grad loop produces
    # the grad_output-weighted VJP. None on every other path.
    loss_grad_output: torch.Tensor | None

    # The optional "when=" buffers (weight_grad_chunk, logits_acc_buf,
    # input_grad_acc_buf, input_chunk_acc_buf, grad_input,
    # grad_linear_weight, grad_linear_bias) are cached_properties below.
    logits_buf: torch.Tensor
    tmp: torch.Tensor
    output: torch.Tensor

    compute_input_grad: bool
    compute_linear_weight_grad: bool
    compute_linear_bias_grad: bool
    alloc_weight_grad_chunk: bool
    alloc_linear_bias_grad_chunk: bool
    alloc_input_grad_acc_buf: bool
    alloc_input_chunk_acc_buf: bool
    alloc_logits_acc_buf: bool
    forward_uses_acc_input: bool
    forward_uses_cuda_out_dtype: bool
    weight_grad_mm_same_dtype: bool
    input_grad_uses_logits_lw: bool
    loop_caches_logits_downcast: bool
    weight_grad_uses_logits_buf_temp: bool

    @cached_property
    def _mask(self) -> torch.Tensor:
        return self.target == self.ignore_index

    @cached_property
    def corrected_target(self) -> torch.Tensor:
        # Replace out-of-range ignore_index values with 0 lazily so
        # downstream index_select / index_add_ stay in bounds.
        if self.ignore_index < 0 or self.ignore_index >= self.num_classes:
            return torch.where(self._mask, 0, self.target)
        return self.target

    @cached_property
    def neg_weight_target(self) -> torch.Tensor:
        # Per-row weighting: -(weight[target] if weight else 1) / d on
        # unmasked positions, 0 on masked positions. d = sum of unmasked
        # weights for "mean", 1 for "sum".
        mask = self._mask
        target = self.corrected_target
        weight = self.weight
        weight_chunk_dtype = self.weight_chunk_dtype
        if weight is None:
            neg_weight_target = (~mask).to(weight_chunk_dtype)
        elif target.numel() > weight.numel():
            neg_weight_target = torch.where(
                mask, 0, weight.to(weight_chunk_dtype).index_select(0, target)
            )
        else:
            neg_weight_target = torch.where(
                mask, 0, weight.index_select(0, target).to(weight_chunk_dtype)
            )
        if self.reduction == "mean":
            d = neg_weight_target.sum()
            neg_weight_target.div_(torch.where(d == 0, torch.nan, -d))
        elif self.reduction == "sum":
            neg_weight_target.neg_()
        else:  # "none"
            # Forward (loss_grad_output is None) intentionally leaves the
            # bare +W[T] unsigned: the loss-only branch needs the positive
            # per-row loss weight. Only the backward (upstream grad set)
            # negates and folds in grad_output[n] for the VJP -- do NOT add
            # a .neg_() here to mirror the "sum" case.
            if self.loss_grad_output is not None:
                neg_weight_target.neg_().mul_(
                    self.loss_grad_output.to(neg_weight_target.dtype)
                )
        return neg_weight_target

    @cached_property
    def linear_weight_cast(self) -> torch.Tensor:
        if self.linear_weight_cast_dtype != self.dtype:
            return self.linear_weight.to(self.linear_weight_cast_dtype)
        return self.linear_weight

    @cached_property
    def linear_bias_cast(self) -> torch.Tensor | None:
        # ``linear_bias`` is added to per-chunk logits, so its dtype
        # must match the logits buffer (which may be acc_dtype on the
        # upcast path). ``None`` when the user passed no bias -- the
        # loop skips the add.
        if self.linear_bias is None:
            return None
        if self.linear_bias.dtype != self.logits_buf.dtype:
            return self.linear_bias.to(self.logits_buf.dtype)
        return self.linear_bias

    @cached_property
    def weight_grad_chunk(self) -> torch.Tensor:
        # (V, F) acc_dtype accumulator for keep-path weight-grad mm;
        # rank-empty on the direct (compact) path.
        return _make_empty(
            (self.num_classes, self.in_features),
            self.acc_dtype,
            self.input.device,
            when=self.alloc_weight_grad_chunk,
        )

    @cached_property
    def logits_acc_buf(self) -> torch.Tensor:
        # (B, V) scratch for non-CUDA mixed-dtype logits_upcast.
        return _make_empty(
            (self.logits_buf.shape[0], self.num_classes),
            self.acc_dtype,
            self.input.device,
            when=self.alloc_logits_acc_buf,
        )

    @cached_property
    def input_grad_acc_buf(self) -> torch.Tensor:
        # (B, F) scratch for input-grad addmm on the buf-and-copy path
        # (CPU mixed-dtype, MPS).
        return _make_empty(
            (self.logits_buf.shape[0], self.in_features),
            self.linear_weight_cast_dtype,
            self.input.device,
            when=self.alloc_input_grad_acc_buf,
        )

    @cached_property
    def input_chunk_acc_buf(self) -> torch.Tensor:
        return _make_empty(
            (self.logits_buf.shape[0], self.in_features),
            self.acc_dtype,
            self.input.device,
            when=self.alloc_input_chunk_acc_buf,
        )

    @cached_property
    def grad_input(self) -> torch.Tensor:
        return _make_empty(
            self.input.shape,
            self.grad_input_dtype,
            self.input.device,
            when=self.compute_input_grad,
        )

    @cached_property
    def grad_linear_weight(self) -> torch.Tensor:
        return _make_zeros(
            self.linear_weight.shape,
            self.dtype,
            self.input.device,
            when=self.compute_linear_weight_grad,
        )

    @cached_property
    def grad_linear_bias(self) -> torch.Tensor:
        # Same shape invariant as ``linear_bias``:
        # ``linear_bias.shape == linear_weight.shape[:-1]``. Final
        # accumulator in input dtype to match ``grad_linear_weight``.
        return _make_zeros(
            self.linear_weight.shape[:-1],
            self.dtype,
            self.input.device,
            when=self.compute_linear_bias_grad,
        )

    @cached_property
    def linear_bias_grad_chunk(self) -> torch.Tensor:
        # ``linear_weight.shape[:-1]`` acc_dtype per-chunk scratch,
        # mirroring ``weight_grad_chunk``'s role for
        # ``grad_linear_weight``: the sub_ + index_add_ pair
        # accumulates per-chunk in acc_dtype, then a single downcast
        # commit (in ``chunks()`` post-yield) lands in
        # ``grad_linear_bias``. ``_make_zeros`` so the first chunk's
        # accumulators start at zero; the post-yield hook zeroes it
        # again for subsequent chunks.
        return _make_zeros(
            self.linear_weight.shape[:-1],
            self.acc_dtype,
            self.input.device,
            when=self.alloc_linear_bias_grad_chunk,
        )

    @classmethod
    def build(
        cls,
        input: torch.Tensor,
        linear_weight: torch.Tensor,
        target: torch.Tensor,
        linear_bias: torch.Tensor | None,
        weight: torch.Tensor | None,
        reduction: str,
        ignore_index: int,
        label_smoothing: float,
        batch_chunk_size: int,
        acc_policy: str,
        acc_dtype: torch.dtype,
        compute_input_grad: bool,
        compute_linear_weight_grad: bool,
        compute_linear_bias_grad: bool,
        loss_grad_output: torch.Tensor | None = None,
    ) -> "_ChunkContext":
        # ===== Validation =====
        if target.dtype != torch.int64:
            raise TypeError(
                f"linear_cross_entropy: target dtype must be torch.int64, got {target.dtype}."
            )
        if label_smoothing > 0.0:
            raise NotImplementedError(
                "linear_cross_entropy does not support label smoothing"
            )
        if reduction not in {"mean", "sum", "none"}:
            raise NotImplementedError(
                f"linear_cross_entropy does not support {reduction=}"
            )
        if linear_bias is not None and linear_bias.shape != linear_weight.shape[:-1]:
            raise RuntimeError(
                "linear_cross_entropy: expected linear_bias shape "
                f"{tuple(linear_weight.shape[:-1])}, got {tuple(linear_bias.shape)}."
            )

        device = input.device
        dtype = input.dtype
        num_batches, in_features = input.shape
        num_classes, _ = linear_weight.shape
        # CUDA gates the out_dtype= mm fast path; the non-CUDA path
        # (CPU, MPS, XPU, ...) routes mixed-dtype mm through explicit
        # casts until out_dtype= is validated on those backends.
        is_cuda = device.type == "cuda"
        is_mps = device.type == "mps"

        if dtype != acc_dtype and not (
            dtype in {torch.float16, torch.bfloat16} and acc_dtype == torch.float32
        ):
            raise RuntimeError(
                "linear_cross_entropy supports float32 acc_dtype with"
                f" float16/bfloat16 inputs, but got {acc_dtype} acc_dtype and {dtype} inputs."
            )
        use_acc_dtype = dtype != acc_dtype

        # Internal dtype layout. ``compact`` reuses ``balanced``'s layout;
        # its savings come from skipping the weight_grad_chunk scratch, not dtype.
        is_memory_like = acc_policy in {"balanced", "compact"}
        if use_acc_dtype:
            output_dtype = acc_dtype if dtype == torch.float16 else dtype
            grad_input_dtype = dtype if is_memory_like else acc_dtype
            # fp16 + memory-like keeps logits at fp16 (2BC vs 4BC);
            # everything else upcasts for softmax stability.
            logits_buf_dtype = (
                dtype if dtype == torch.float16 and is_memory_like else acc_dtype
            )
            # acc_dtype mandatory: 1/N for reduction="mean" is subnormal
            # in fp16 at N >= 65536.
            weight_chunk_dtype = acc_dtype
        else:
            output_dtype = grad_input_dtype = logits_buf_dtype = weight_chunk_dtype = (
                dtype
            )

        # ===== Dispatch flags =====
        # CUDA + balanced uses cuBLAS out_dtype= directly; no cast.
        needs_linear_weight_cast = use_acc_dtype and (
            not is_cuda or (compute_input_grad and grad_input_dtype == logits_buf_dtype)
        )
        linear_weight_cast_dtype = (
            logits_buf_dtype if needs_linear_weight_cast else dtype
        )
        alloc_weight_grad_chunk = compute_linear_weight_grad and not (
            acc_policy == "compact" and (is_cuda or logits_buf_dtype == dtype)
        )
        # Per-chunk acc_dtype scratch for grad_linear_bias; same
        # bulk-+-correction precision rationale as
        # ``alloc_weight_grad_chunk``. Fires only when there is
        # precision to gain: ``use_acc_dtype`` and the user asked for
        # the bias gradient.
        alloc_linear_bias_grad_chunk = compute_linear_bias_grad and use_acc_dtype
        alloc_input_grad_acc_buf = (
            compute_input_grad
            and not is_cuda
            and (grad_input_dtype != linear_weight_cast_dtype or is_mps)
        )
        alloc_input_chunk_acc_buf = use_acc_dtype and (
            compute_linear_weight_grad or (not is_cuda and dtype != logits_buf_dtype)
        )
        forward_uses_acc_input = (
            use_acc_dtype and not is_cuda and dtype != logits_buf_dtype
        )
        forward_uses_cuda_out_dtype = is_cuda and use_acc_dtype
        weight_grad_mm_same_dtype = logits_buf_dtype == acc_dtype
        # Storage-trick fires when logits' dtype differs from the input-
        # grad accumulator dtype.
        input_grad_uses_logits_lw = (
            is_cuda and compute_input_grad and logits_buf_dtype != grad_input_dtype
        )
        # Per-iter ``logits.to(linear_weight.dtype)`` shared between the
        # inlined input-grad addmm and the inlined direct weight-grad addmm_.
        loop_caches_logits_downcast = input_grad_uses_logits_lw or (
            compute_linear_weight_grad
            and not alloc_weight_grad_chunk
            and logits_buf_dtype != dtype
        )
        alloc_logits_acc_buf = (
            alloc_weight_grad_chunk and not is_cuda and logits_buf_dtype != acc_dtype
        )
        # Direct weight-grad path: reuse logits_buf storage for the
        # acc_dtype->dtype cast when its bytes/row fits the (B, F) cast.
        weight_grad_uses_logits_buf_temp = (
            not alloc_weight_grad_chunk
            and use_acc_dtype
            and num_classes * logits_buf_dtype.itemsize >= in_features * dtype.itemsize
        )

        return cls(
            dtype=dtype,
            num_batches=num_batches,
            in_features=in_features,
            num_classes=num_classes,
            is_cuda=is_cuda,
            use_acc_dtype=use_acc_dtype,
            acc_dtype=acc_dtype,
            weight_chunk_dtype=weight_chunk_dtype,
            grad_input_dtype=grad_input_dtype,
            linear_weight_cast_dtype=linear_weight_cast_dtype,
            input=input,
            target=target,
            weight=weight,
            ignore_index=ignore_index,
            reduction=reduction,
            linear_weight=linear_weight,
            linear_bias=linear_bias,
            loss_grad_output=loss_grad_output,
            logits_buf=torch.empty(
                (batch_chunk_size, num_classes),
                dtype=logits_buf_dtype,
                device=device,
                requires_grad=False,
            ),
            tmp=torch.empty(
                (batch_chunk_size,),
                dtype=logits_buf_dtype,
                device=device,
                requires_grad=False,
            ),
            output=_make_zeros((), output_dtype, device),
            compute_input_grad=compute_input_grad,
            compute_linear_weight_grad=compute_linear_weight_grad,
            compute_linear_bias_grad=compute_linear_bias_grad,
            alloc_weight_grad_chunk=alloc_weight_grad_chunk,
            alloc_linear_bias_grad_chunk=alloc_linear_bias_grad_chunk,
            alloc_input_grad_acc_buf=alloc_input_grad_acc_buf,
            alloc_input_chunk_acc_buf=alloc_input_chunk_acc_buf,
            alloc_logits_acc_buf=alloc_logits_acc_buf,
            forward_uses_acc_input=forward_uses_acc_input,
            forward_uses_cuda_out_dtype=forward_uses_cuda_out_dtype,
            weight_grad_mm_same_dtype=weight_grad_mm_same_dtype,
            input_grad_uses_logits_lw=input_grad_uses_logits_lw,
            loop_caches_logits_downcast=loop_caches_logits_downcast,
            weight_grad_uses_logits_buf_temp=weight_grad_uses_logits_buf_temp,
        )

    def chunks(self):
        """Yield a ``_ChunkViews`` per iter. Post-yield, run the
        scratch-to-final commits that the loop body deferred: the
        buf-and-copy input-grad slice into ``grad_input`` (skipped on
        the fast path where ``grad_input_chunk`` aliases
        ``grad_input``), and the acc_dtype linear_bias scratch into
        ``grad_linear_bias`` (skipped when accumulation went straight
        into ``grad_linear_bias`` in input dtype).
        """
        batch_chunk_size = self.logits_buf.shape[0]
        for bchunk_start in range(0, self.num_batches, batch_chunk_size):
            bchunk_size = min(batch_chunk_size, self.num_batches - bchunk_start)
            chunk = self.bind_chunk(bchunk_start, bchunk_size)
            yield chunk
            if self.alloc_input_grad_acc_buf:
                self.grad_input.narrow(0, bchunk_start, bchunk_size).copy_(
                    chunk.grad_input_chunk
                )
            if self.alloc_linear_bias_grad_chunk:
                # ``chunk.bias_grad_acc`` aliases ``linear_bias_grad_chunk``
                # here; downcast-commit, then zero for the next chunk.
                self.grad_linear_bias.add_(
                    chunk.bias_grad_acc.to(self.grad_linear_bias.dtype)
                )
                chunk.bias_grad_acc.zero_()

    def bind_chunk(self, bchunk_start: int, bchunk_size: int) -> _ChunkViews:
        input_chunk = self.input.narrow(0, bchunk_start, bchunk_size)
        target_chunk = self.corrected_target.narrow(0, bchunk_start, bchunk_size)
        weight_chunk = self.neg_weight_target.narrow(0, bchunk_start, bchunk_size)
        logits = self.logits_buf.narrow(0, 0, bchunk_size)
        input_chunk_acc = (
            self.input_chunk_acc_buf.narrow(0, 0, bchunk_size).copy_(input_chunk)
            if self.alloc_input_chunk_acc_buf
            else input_chunk
        )
        return _ChunkViews(
            ctx=self,
            bchunk_start=bchunk_start,
            bchunk_size=bchunk_size,
            input_chunk=input_chunk,
            target_chunk=target_chunk,
            weight_chunk=weight_chunk,
            logits=logits,
            input_chunk_acc=input_chunk_acc,
        )

    def mm(self, mat1: torch.Tensor, mat2: torch.Tensor, *, out: torch.Tensor) -> None:
        # Mismatched dtype: cuBLAS ``out_dtype=`` upcasts inside the matmul.
        if mat1.dtype == out.dtype:
            torch.mm(mat1, mat2, out=out)
        else:
            torch.mm(mat1, mat2, out_dtype=out.dtype, out=out)

    def amax(self, x: torch.Tensor) -> torch.Tensor:
        out = self.tmp.narrow(0, 0, x.shape[0]).unsqueeze(1)
        torch.amax(x, dim=1, keepdim=True, out=out)
        return out

    def dotgather(
        self, weight: torch.Tensor, x: torch.Tensor, indices: torch.Tensor
    ) -> torch.Tensor:
        return weight.dot(x.gather(1, indices.unsqueeze(1)).squeeze(1).to(weight.dtype))

    def sumexp_(self, x: torch.Tensor, dim: int) -> torch.Tensor:
        # Result always in acc_dtype: fp16 softmax_denom can underflow
        # on rows with widely-spread logits, after which log(0)=-inf
        # poisons the loss and 1/0 the gradient.
        return x.exp_().sum(dim, dtype=self.acc_dtype)

    def sum(self, x: torch.Tensor, dim: int) -> torch.Tensor:
        # Reduce ``x`` along ``dim`` at acc_dtype on mixed-precision
        # paths and at input dtype otherwise -- the precision
        # contract callers see is "as accurate as the active policy
        # allows". Currently used by the bias-grad path so the per-
        # chunk sum carries full precision before the post-yield
        # downcast commit to grad_linear_bias.
        target_dtype = self.acc_dtype if self.use_acc_dtype else self.dtype
        return x.sum(dim, dtype=target_dtype)

    def div(self, num: torch.Tensor, den: torch.Tensor) -> torch.Tensor:
        factor = self.tmp.narrow(0, 0, num.shape[0])
        # MPS lacks fused cross-dtype torch.div(out=); copy_ fallback.
        if self.input.device.type == "mps" and (
            num.dtype != factor.dtype or den.dtype != factor.dtype
        ):
            factor.copy_(num / den)
        else:
            torch.div(num, den, out=factor)
        return factor

    def mul(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
        # Pick a destination from in-flight scratch when possible.
        if self.use_acc_dtype:
            return x.mul_(w.unsqueeze(1))
        if self.num_classes >= self.in_features:
            return torch.mul(
                x,
                w.unsqueeze(1),
                out=self.logits_buf.narrow(0, 0, x.shape[0]).narrow(
                    1, 0, self.in_features
                ),
            )
        return x * w.unsqueeze(1)

    def to(self, x: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor:
        # Reuse logits_buf storage as cast destination when its bytes/row fits.
        if x.dtype == dtype:
            return x
        if self.weight_grad_uses_logits_buf_temp:
            x = (
                self.logits_buf.view(dtype)
                .narrow(0, 0, x.shape[0])
                .narrow(1, 0, x.shape[1])
                .copy_(x)
            )
            return x
        return x.to(dtype)


def _linear_cross_entropy_batch_chunked_accumulator(
    input: torch.Tensor,
    linear_weight: torch.Tensor,
    target: torch.Tensor,
    linear_bias: torch.Tensor | None,
    weight: torch.Tensor | None,
    reduction: str,
    ignore_index: int,
    label_smoothing: float,
    batch_chunk_size: int,
    acc_policy: str,
    acc_dtype: torch.dtype,
    compute_input_grad: bool,
    compute_linear_weight_grad: bool,
    compute_linear_bias_grad: bool,
    loss_grad_output: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Chunked loop shared by all three custom-op entry points; returns
    ``(loss, grad_input, grad_linear_weight, grad_linear_bias)``.

    Caller -> branch:
      scalar-op forward (reduction in {mean,sum})  -> main grad loop
      no_reduction forward (none, loss_grad_output None) -> loss-only branch
      no_reduction backward (none, loss_grad_output set) -> main grad loop

    - ``_linear_cross_entropy_batch_chunked`` (scalar reduction): the
      forward precomputes the scalar loss and all gradients in one pass;
      backward just scales the stashed grads by the upstream gradient.
    - ``_linear_cross_entropy_batch_chunked_no_reduction`` (forward):
      ``reduction='none'`` with ``loss_grad_output is None`` takes the
      loss-only branch, returning the per-sample loss ``(N,)`` and no
      gradients (they are recomputed in its backward instead).
    - ``_linear_cross_entropy_batch_chunked_no_reduction`` (backward):
      ``loss_grad_output`` set folds the per-sample upstream grad into
      ``neg_weight_target``, so this same loop recomputes the chunked
      grads as the grad_output-weighted VJP (the scalar loss it also
      accumulates is then meaningless and ignored).

    Dispatch contract
    -----------------
    The chunked loop body is the mathematical description of the
    algorithm. All policy / device / dtype dispatch lives in
    ``_ChunkContext`` (and its ``_ChunkViews`` peers):

    - **Buffer choices** -- whether a per-chunk acc_dtype scratch is
      used or accumulation runs directly into the final tensor -- live
      in ``alloc_*`` flags computed in ``_ChunkContext.build`` and
      surfaced as ``_ChunkViews`` properties (e.g.
      ``grad_input_chunk``, ``bias_grad_acc``).
    - **Scratch-to-final commits** for those buffers live in
      ``_ChunkContext.chunks()`` post-yield hooks (e.g. the
      ``grad_input`` buf-and-copy commit, the ``grad_linear_bias``
      downcast-and-zero commit).
    - **Dtype dispatch on individual math operations** lives in
      ``_ChunkContext`` methods that mirror PyTorch ops --
      ``mm``, ``amax``, ``dotgather``, ``sumexp_``, ``sum``, ``div``,
      ``mul``, ``to`` -- each one hiding ``out_dtype=`` / acc-dtype /
      buffer-reuse decisions behind a single math call.

    The function body therefore should not introduce inline
    ``if ctx.use_acc_dtype`` / ``if ctx.is_cuda`` / etc. branches; new
    policy-aware behaviour belongs in one of the three locations
    above. The two existing inline branches (``if compute_input_grad``,
    ``if compute_linear_weight_grad``, etc.) gate optional outputs,
    not dispatch.
    """
    # Direct callers must resolve "auto" / None via _adjust first.
    if acc_policy == "auto" or acc_dtype is None:
        raise RuntimeError(
            f"unresolved acc_policy={acc_policy!r} or acc_dtype={acc_dtype!r};"
            " use F.linear_cross_entropy."
        )
    ctx = _ChunkContext.build(
        input,
        linear_weight,
        target,
        linear_bias,
        weight,
        reduction,
        ignore_index,
        label_smoothing,
        batch_chunk_size,
        acc_policy,
        acc_dtype,
        compute_input_grad,
        compute_linear_weight_grad,
        compute_linear_bias_grad,
        loss_grad_output=loss_grad_output,
    )
    dtype = ctx.dtype
    # reduction='none' forward (no upstream grad yet): per-sample loss
    # into an (N,) output, no gradient precompute. ``chunk.weight_chunk``
    # is the unsigned masked class weight here (neg_weight_target's "none"
    # forward form), so each row's loss is W[T[n]] * (log denom -
    # shifted_logit[T[n]]) = W[T[n]] * (-log_softmax). The "none" backward
    # (loss_grad_output set) routes through the grad loop below instead.
    if reduction == "none" and loss_grad_output is None:
        linear_bias_cast = ctx.linear_bias_cast
        out = torch.empty(ctx.num_batches, dtype=dtype, device=ctx.input.device)
        for chunk in ctx.chunks():
            logits = chunk.logits
            ctx.mm(chunk.input, chunk.linear_weight.T, out=logits)
            if linear_bias_cast is not None:
                logits.add_(linear_bias_cast)
            logits.sub_(ctx.amax(logits))
            # Read the target logit BEFORE ``sumexp_`` -- it does ``exp_()``
            # in place, overwriting logits with exp(shifted).
            ls_target = logits.gather(1, chunk.target_chunk.unsqueeze(1)).squeeze(1)
            softmax_denom = ctx.sumexp_(logits, dim=1)
            loss_chunk = softmax_denom.log_().sub_(ls_target.to(softmax_denom.dtype))
            loss_chunk.mul_(chunk.weight_chunk.to(softmax_denom.dtype))
            out.narrow(0, chunk.bchunk_start, chunk.bchunk_size).copy_(loss_chunk)
        return (
            out,
            ctx.grad_input.to(dtype),
            ctx.grad_linear_weight,
            ctx.grad_linear_bias,
        )
    output = ctx.output
    grad_input = ctx.grad_input
    grad_linear_weight = ctx.grad_linear_weight
    grad_linear_bias = ctx.grad_linear_bias
    weight_grad_chunk = ctx.weight_grad_chunk
    linear_weight_cast = ctx.linear_weight_cast
    linear_bias_cast = ctx.linear_bias_cast

    if reduction == "mean" and ctx.num_batches == 0:
        output.fill_(torch.nan)

    compute_grads = (
        compute_input_grad or compute_linear_weight_grad or compute_linear_bias_grad
    )

    # Do not ``break`` from this loop -- ``ctx.chunks()`` runs a
    # post-yield grad_input commit that the in-flight chunk needs.
    for chunk in ctx.chunks():
        logits = chunk.logits
        weight_chunk = chunk.weight_chunk
        target_chunk = chunk.target_chunk

        ctx.mm(chunk.input, chunk.linear_weight.T, out=logits)
        if linear_bias_cast is not None:
            logits.add_(linear_bias_cast)  # broadcasts (V,) -> (B_chunk, V)
        logits.sub_(ctx.amax(logits))  # softmax stability
        # output += <weight_chunk, logits[:, target_chunk]>
        output.add_(ctx.dotgather(weight_chunk, logits, target_chunk))
        softmax_denom = ctx.sumexp_(logits, dim=1)
        if compute_grads:
            # logits *= weight_chunk / softmax_denom         (= softmax(logits) * w)
            # MPS: in-place mul_ on a narrow view of logits_buf doesn't
            # propagate to parent storage; the out= form does.
            torch.mul(
                logits,
                ctx.div(weight_chunk, softmax_denom).unsqueeze(1),
                out=logits,
            )
        # output -= <weight_chunk, log(softmax_denom)>
        # softmax_denom is in acc_dtype (see ``sumexp_``); promote weight_chunk
        # for the dot rather than rounding the wider denominator down.
        output.sub_(weight_chunk.to(softmax_denom.dtype).dot(softmax_denom.log_()))

        if compute_grads:
            if compute_linear_bias_grad:
                # Must run BEFORE the linear-weight-grad branch below.
                # The ordering invariant is defensive across all
                # ``ctx.mul`` paths: only the (not use_acc_dtype &&
                # num_classes >= in_features) branch overwrites the
                # first ``in_features`` columns of ``logits_buf`` via a
                # narrow ``out=`` view today, but ``ctx.mul``'s contract
                # does not promise to leave ``logits_buf`` untouched on
                # other paths -- placing the bias-grad read here keeps
                # the loop correct regardless of which ``ctx.mul``
                # variant fires.
                #
                # The inline form expresses
                #   dL/db_j = sum_n W[T[n]] * (softmax(X)[n, j] - delta(T[n], j))
                # as ``-logits.sum(dim=0)`` (the softmax-weighted term)
                # plus ``index_add_(0, target_chunk, weight_chunk)``
                # (the one-hot-weighted correction). ``logits`` is
                # ``softmax(X) * weight_chunk[n]`` here; ``weight_chunk``
                # already carries the sign and the 1/d mean
                # normalization (``W`` here is the CE class weight, not
                # ``linear_weight``), so no extra scaling is needed.
                bias_grad_acc = chunk.bias_grad_acc
                bias_grad_acc.sub_(ctx.sum(logits, dim=0))
                bias_grad_acc.index_add_(0, target_chunk, weight_chunk)
            if compute_input_grad:
                grad_input_chunk = chunk.grad_input_chunk
                input_grad_logits = chunk.input_grad_logits
                input_grad_linear_weight = chunk.input_grad_linear_weight

                # grad_input_chunk = linear_weight_cast[target_chunk] * weight_chunk
                torch.mul(
                    torch.index_select(linear_weight_cast, 0, target_chunk),
                    weight_chunk.to(grad_input_chunk.dtype).unsqueeze(1),
                    out=grad_input_chunk,
                )
                # grad_input_chunk -= input_grad_logits @ input_grad_linear_weight
                torch.addmm(
                    grad_input_chunk,
                    input_grad_logits,
                    input_grad_linear_weight,
                    alpha=-1,
                    out=grad_input_chunk,
                )
            if compute_linear_weight_grad:
                input_chunk_acc = chunk.input_chunk_acc

                # grad_linear_weight += per-chunk weight grad
                # (= -logits.T @ input + input*weight scattered at target rows)
                if ctx.alloc_weight_grad_chunk:
                    # Stage per-chunk weight grad in acc_dtype scratch;
                    # one sub_ commit keeps bulk-minus-correction in fp32.
                    logits_upcast = chunk.logits_upcast
                    weight_grad_input = chunk.weight_grad_input
                    ctx.mm(
                        logits_upcast.T,
                        weight_grad_input,
                        out=weight_grad_chunk,
                    )
                    temp = ctx.mul(input_chunk_acc, weight_chunk)
                    weight_grad_chunk.index_add_(0, target_chunk, temp, alpha=-1)
                    grad_linear_weight.sub_(weight_grad_chunk)
                else:
                    # Stream bulk + correction directly into grad_lw without
                    # a (V, F) scratch (the compact path).
                    grad_linear_weight.addmm_(
                        chunk.logits_downcast.T, chunk.input_chunk, alpha=-1
                    )
                    temp = ctx.to(
                        ctx.mul(input_chunk_acc, weight_chunk),
                        dtype=grad_linear_weight.dtype,
                    )
                    grad_linear_weight.index_add_(0, target_chunk, temp, alpha=1)

    return (
        output.to(dtype),
        grad_input.to(dtype),
        grad_linear_weight,
        grad_linear_bias,
    )


# Scalar-reduction (mean/sum) op. Returns
# ``(loss, grad_input, grad_linear_weight, grad_linear_bias)`` with grads
# precomputed in forward (the accumulator runs the full grad loop) and
# stashed on ctx by ``setup_context``; backward mutates them in place via
# ``.mul_()`` by the upstream gradient. ``torch.library.register_autograd``
# has no ``save_for_backward``-style state, hence the return-tuple +
# ctx-stash dance. ``mutates_args=()`` reflects that inputs are not mutated.
# Shape invariants: grad_input.shape == input.shape,
# grad_linear_weight.shape == linear_weight.shape, grad_linear_bias.shape ==
# linear_weight.shape[:-1] (here (C,), since the upstream gate requires
# out_features == ()).
@torch.library.custom_op(
    "torch_nn::_linear_cross_entropy_batch_chunked", mutates_args=()
)
def _linear_cross_entropy_batch_chunked(
    input: torch.Tensor,
    linear_weight: torch.Tensor,
    target: torch.Tensor,
    linear_bias: torch.Tensor | None,
    weight: torch.Tensor | None,
    reduction: str,
    ignore_index: int,
    label_smoothing: float,
    batch_chunk_size: int,
    acc_policy: str,
    acc_dtype: torch.dtype,
    allow_retain_graph: bool,
    compute_input_grad: bool,
    compute_linear_weight_grad: bool,
    compute_linear_bias_grad: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Scalar-reduction (mean/sum) chunked op: precomputes the loss and
    gradients in forward, backward scales by the upstream gradient. The
    chunked math lives in ``_linear_cross_entropy_batch_chunked_accumulator``.
    """
    # AOTAutograd/AOTInductor bake compute_*_grad at trace time; catch
    # the silent-corruption case (False at trace, but grad-enabled at
    # runtime with a requires_grad leaf). These guards are specific to
    # the precompute path -- the reduction='none' op recomputes grads in
    # its own backward, so it does not route through here.
    grad_enabled = torch.is_grad_enabled()
    if not compute_input_grad and input.requires_grad and grad_enabled:
        raise RuntimeError(
            "linear_cross_entropy chunked op: compute_input_grad was False at "
            "trace time but input.requires_grad is True at runtime; recompile "
            "the graph with the desired requires_grad."
        )
    if not compute_linear_weight_grad and linear_weight.requires_grad and grad_enabled:
        raise RuntimeError(
            "linear_cross_entropy chunked op: compute_linear_weight_grad was "
            "False at trace time but linear_weight.requires_grad is True at "
            "runtime; recompile the graph with the desired requires_grad."
        )
    if (
        not compute_linear_bias_grad
        and linear_bias is not None
        and linear_bias.requires_grad
        and grad_enabled
    ):
        raise RuntimeError(
            "linear_cross_entropy chunked op: compute_linear_bias_grad was False at "
            "trace time but linear_bias.requires_grad is True at runtime; recompile "
            "the graph with the desired requires_grad."
        )
    return _linear_cross_entropy_batch_chunked_accumulator(
        input,
        linear_weight,
        target,
        linear_bias,
        weight,
        reduction,
        ignore_index,
        label_smoothing,
        batch_chunk_size,
        acc_policy,
        acc_dtype,
        compute_input_grad,
        compute_linear_weight_grad,
        compute_linear_bias_grad,
    )


# Op arg count; backward returns this many slots. Hand-maintained
# because CustomOpDef has no public arg-introspection. Update on
# signature change.
_NUM_OP_INPUTS = 15


@_linear_cross_entropy_batch_chunked.register_fake
def _(
    input,
    linear_weight,
    target,
    linear_bias,
    weight,
    reduction,
    ignore_index,
    label_smoothing,
    batch_chunk_size,
    acc_policy: str,
    acc_dtype,
    allow_retain_graph,
    compute_input_grad,
    compute_linear_weight_grad,
    compute_linear_bias_grad,
):
    if reduction in {"mean", "sum"}:
        result = torch.empty((), dtype=input.dtype, device=input.device)
    else:
        raise NotImplementedError(f"linear_cross_entropy does not support {reduction=}")
    grad_input_shape = input.shape if compute_input_grad else (0, 0)
    grad_linear_weight_shape = (
        linear_weight.shape if compute_linear_weight_grad else (0, 0)
    )
    # ``linear_bias.shape == linear_weight.shape[:-1]`` (the (C, *out_features)
    # invariant); on this code path ``out_features == ()`` so it reduces
    # to (C,), but expressing it this way keeps the relationship visible.
    grad_linear_bias_shape = (
        linear_weight.shape[:-1] if compute_linear_bias_grad else (0,)
    )

    grad_input = torch.empty(
        grad_input_shape,
        dtype=input.dtype,
        device=input.device,
        requires_grad=False,
    )
    # The real implementation allocates ``grad_linear_weight`` and
    # ``grad_linear_bias`` in ``self.dtype = input.dtype`` (chunked op
    # supports mixed-precision: ``input.dtype`` may differ from
    # ``linear_weight.dtype``). Mirror that here so the fake-tensor
    # shape/dtype matches what eager produces.
    grad_linear_weight = torch.empty(
        grad_linear_weight_shape,
        dtype=input.dtype,
        device=linear_weight.device,
        requires_grad=False,
    )
    grad_linear_bias = torch.empty(
        grad_linear_bias_shape,
        dtype=input.dtype,
        device=input.device,
        requires_grad=False,
    )
    return result, grad_input, grad_linear_weight, grad_linear_bias


@_linear_cross_entropy_batch_chunked.register_vmap
def _vmap(info, in_dims, *args):
    """vmap rule (slow path: per-sample Python loop). A fold-into-num_batches
    fast path is now possible for both ops via the reduction='none' chunked op
    (``_linear_cross_entropy_batch_chunked_no_reduction``); left as a follow-up.
    """
    batch_size = info.batch_size

    moved_args = [
        arg.movedim(in_dim, 0)
        if in_dim is not None and isinstance(arg, torch.Tensor)
        else arg
        for arg, in_dim in zip(args, in_dims)
    ]

    outputs = [
        _linear_cross_entropy_batch_chunked(
            *[
                arg[i] if in_dim is not None and isinstance(arg, torch.Tensor) else arg
                for arg, in_dim in zip(moved_args, in_dims)
            ]
        )
        for i in range(batch_size)
    ]
    losses = torch.stack([o[0] for o in outputs])
    grad_inputs = torch.stack([o[1] for o in outputs])
    grad_linear_weights = torch.stack([o[2] for o in outputs])
    grad_linear_biases = torch.stack([o[3] for o in outputs])
    return (losses, grad_inputs, grad_linear_weights, grad_linear_biases), (0, 0, 0, 0)


def _linear_cross_entropy_batch_chunked_backward(
    ctx, grad_output, _gi_grad, _gw_grad, _gb_grad
):
    result = [None] * _NUM_OP_INPUTS
    if ctx.compute_input_grad:
        if ctx.allow_retain_graph:
            result[0] = ctx._gi * grad_output
        else:
            gi, ctx._gi = ctx._gi, None
            if gi is None:
                raise RuntimeError(
                    "linear_cross_entropy chunked backward called twice; "
                    "retain_graph=True / double backward are not supported."
                )
            result[0] = gi.mul_(grad_output)

    if ctx.compute_linear_weight_grad:
        if ctx.allow_retain_graph:
            result[1] = ctx._gw * grad_output
        else:
            gw, ctx._gw = ctx._gw, None
            if gw is None:
                raise RuntimeError(
                    "linear_cross_entropy chunked backward called twice; "
                    "retain_graph=True / double backward are not supported."
                )
            result[1] = gw.mul_(grad_output)

    # ``linear_bias`` is forward arg position 3 (after input, linear_weight, target).
    if ctx.compute_linear_bias_grad:
        if ctx.allow_retain_graph:
            result[3] = ctx._gb * grad_output
        else:
            gb, ctx._gb = ctx._gb, None
            if gb is None:
                raise RuntimeError(
                    "linear_cross_entropy chunked backward called twice; "
                    "retain_graph=True / double backward are not supported."
                )
            result[3] = gb.mul_(grad_output)
    return tuple(result)


_linear_cross_entropy_batch_chunked.register_autograd(
    _linear_cross_entropy_batch_chunked_backward,
    setup_context=_linear_cross_entropy_batch_chunked_setup_context,
)


# reduction='none' op. Separate from the scalar-reduction op because the
# autograd model differs: per-sample loss makes ``grad_output`` a vector,
# so grad_linear_weight / grad_linear_bias cannot be precomputed in forward
# (they need the per-sample upstream grad before summing over the batch, and
# storing the per-sample (N, V) softmax would defeat chunking). Instead the
# forward returns only the loss and saves its inputs; backward recomputes the
# chunked grads with ``grad_output`` folded into ``weight_chunk`` (both share
# ``_linear_cross_entropy_batch_chunked_accumulator``). Peak memory stays bounded to
# (batch_chunk_size, num_classes); the cost is one recomputed logits matmul
# per chunk in backward. Recompute-from-saved-inputs supports retain_graph
# naturally (no consumed buffers), so no allow_retain_graph flag is needed.
@torch.library.custom_op(
    "torch_nn::_linear_cross_entropy_batch_chunked_no_reduction", mutates_args=()
)
def _linear_cross_entropy_batch_chunked_no_reduction(
    input: torch.Tensor,
    linear_weight: torch.Tensor,
    target: torch.Tensor,
    linear_bias: torch.Tensor | None,
    weight: torch.Tensor | None,
    ignore_index: int,
    batch_chunk_size: int,
    acc_policy: str,
    acc_dtype: torch.dtype,
) -> torch.Tensor:
    """reduction='none' chunked forward: returns per-sample loss ``(N,)``.

    Gradients are not precomputed (the per-sample upstream grad is
    unknown at forward); the accumulator's loss-only branch runs here,
    and the backward recomputes the chunked grads. See
    ``_linear_cross_entropy_batch_chunked_accumulator``.
    """
    # loss_grad_output=None -> the accumulator's reduction='none' forward
    # (loss-only) branch; compute_*=False since nothing is precomputed.
    return _linear_cross_entropy_batch_chunked_accumulator(
        input,
        linear_weight,
        target,
        linear_bias,
        weight,
        "none",
        ignore_index,
        0.0,
        batch_chunk_size,
        acc_policy,
        acc_dtype,
        compute_input_grad=False,
        compute_linear_weight_grad=False,
        compute_linear_bias_grad=False,
    )[0]


def _linear_cross_entropy_batch_chunked_no_reduction_setup_context(ctx, inputs, output):
    (
        input,
        linear_weight,
        target,
        linear_bias,
        weight,
        ignore_index,
        batch_chunk_size,
        acc_policy,
        acc_dtype,
    ) = inputs
    # Detach before saving: the backward recomputes the grads analytically
    # from these values (it does not differentiate through them), and saving
    # ``linear_weight`` -- a differentiable ``reshape`` view of the
    # ``nn.Linear`` weight -- would make its SavedVariable retain the base
    # parameter, inflating its use_count past backward and breaking
    # swap-based ``module.to()`` (the custom-op node keeps saved tensors
    # alive while the loss output is still referenced).
    ctx.save_for_backward(
        input.detach(),
        linear_weight.detach(),
        target,
        linear_bias.detach() if linear_bias is not None else None,
        weight.detach() if weight is not None else None,
    )
    ctx._ignore_index = ignore_index
    ctx._batch_chunk_size = batch_chunk_size
    ctx._acc_policy = acc_policy
    ctx._acc_dtype = acc_dtype


@_linear_cross_entropy_batch_chunked_no_reduction.register_fake
def _(
    input,
    linear_weight,
    target,
    linear_bias,
    weight,
    ignore_index,
    batch_chunk_size,
    acc_policy: str,
    acc_dtype,
):
    return torch.empty(
        input.shape[:1], dtype=input.dtype, device=input.device, requires_grad=False
    )


@_linear_cross_entropy_batch_chunked_no_reduction.register_vmap
def _vmap_no_reduction(info, in_dims, *args):
    """vmap rule (slow path: per-sample Python loop, mirrors the
    scalar-reduction op). Stacks the per-sample ``(N,)`` losses.
    """
    batch_size = info.batch_size
    moved_args = [
        arg.movedim(in_dim, 0)
        if in_dim is not None and isinstance(arg, torch.Tensor)
        else arg
        for arg, in_dim in zip(args, in_dims)
    ]
    outputs = [
        _linear_cross_entropy_batch_chunked_no_reduction(
            *[
                arg[i] if in_dim is not None and isinstance(arg, torch.Tensor) else arg
                for arg, in_dim in zip(moved_args, in_dims)
            ]
        )
        for i in range(batch_size)
    ]
    return torch.stack(outputs), 0


def _linear_cross_entropy_batch_chunked_no_reduction_backward(ctx, grad_output):
    # grad_output is the per-sample upstream grad (N,). Forward arg order:
    # (input, linear_weight, target, linear_bias, weight, ignore_index,
    #  batch_chunk_size, acc_policy, acc_dtype).
    input, linear_weight, target, linear_bias, weight = ctx.saved_tensors
    needs = ctx.needs_input_grad
    compute_input_grad = needs[0]
    compute_linear_weight_grad = needs[1]
    compute_linear_bias_grad = linear_bias is not None and needs[3]
    result: list[torch.Tensor | None] = [None] * 9
    if not (
        compute_input_grad or compute_linear_weight_grad or compute_linear_bias_grad
    ):
        return tuple(result)
    # loss_grad_output=grad_output -> the accumulator's grad loop with the
    # per-sample upstream grad folded into neg_weight_target (the "none"
    # backward path), recomputing the chunked grads as the VJP. The loss
    # (result[0]) it also returns is the meaningless scalar and dropped.
    _, grad_input, grad_linear_weight, grad_linear_bias = (
        _linear_cross_entropy_batch_chunked_accumulator(
            input,
            linear_weight,
            target,
            linear_bias,
            weight,
            "none",
            ctx._ignore_index,
            0.0,  # label_smoothing
            ctx._batch_chunk_size,
            ctx._acc_policy,
            ctx._acc_dtype,
            compute_input_grad,
            compute_linear_weight_grad,
            compute_linear_bias_grad,
            loss_grad_output=grad_output,
        )
    )
    if compute_input_grad:
        result[0] = grad_input
    if compute_linear_weight_grad:
        result[1] = grad_linear_weight
    if compute_linear_bias_grad:
        result[3] = grad_linear_bias
    return tuple(result)


_linear_cross_entropy_batch_chunked_no_reduction.register_autograd(
    _linear_cross_entropy_batch_chunked_no_reduction_backward,
    setup_context=_linear_cross_entropy_batch_chunked_no_reduction_setup_context,
)
