TensorBackward.vjp_cross_entropy: def of bendtorch-core@0.1.0.0
Import: import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as TensorBackward
Lines 505-520 of bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend: https://jend.leeeet.dev/read.txt?path=bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend&start_line=505&end_line=520

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/elementwise.bend as Elementwise
import bendtorch-core@0.1.0.0/src/autograd/loss.bend as Loss

# d cross_entropy(a)/da, as two contributions to a that its slot adds in
# order. The first is PyTorch's nll_loss backward, OneHotGrad of g and the
# saved one-hot rows: -(g / N) at each row's target and +0.0 elsewhere, for
# the N rows. The second is the rest of PyTorch's log_softmax backward,
# SoftmaxGrad of g and the saved lp: -(exp(lp) * r) with r the row sum of the
# first, -(g / N) + 0.0, or -(g / N) itself when the logits have one column
# (PyTorch's sum of one word keeps -0.0). Added into an empty slot, first + second is
# PyTorch's first - exp(lp) * r word for word. The steps are the Mul VJP's,
# with OneHotGrad and SoftmaxGrad for Mul and both contributions to a.
def vjp_cross_entropy(+a: Nat, +input: Shape.Shape, saved: Saved, g: Core.Dynamic) -> Pass:
  match saved:
    case Unsaved{}:
      Pass{Unsaved{}, g, Fail{TensorError{Core.StorageLengthMismatch{}}}}
    case Operands{lp, onehot}:
      entropy_first(a, input, Loss.count_f32(Loss.rows(input)), Loss.one_class(input), lp,
        Elementwise.binary(Elementwise.OneHotGrad{Loss.count_f32(Loss.rows(input))}, g, onehot))

Uses (14):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Result
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/loss.bend&name=count_f32
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/loss.bend&name=one_class
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/loss.bend&name=rows
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend&name=Error
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend&name=Pass
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend&name=Saved
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend&name=entropy_first
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/shape/core.bend&name=Shape
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/core.bend&name=Dynamic
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/core.bend&name=Error
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/elementwise.bend&name=Operation
  https://jend.leeeet.dev/definition.txt?path=bendtorch-core@0.1.0.0/src/tensor/elementwise.bend&name=binary