def vjp_cross_entropy

bendtorch-core@0.1.0.0 · src/autograd/tensor_backward.bend:505

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

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