def cross_entropy

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

import Base

1658# Records cross_entropy(a, targets), PyTorch's cross_entropy with class
1659# indices and mean reduction, a rank-0 tensor, from a copy of a's value: a
1660# must be [rows, columns] logits and `targets` one class below `columns`
1661# per row. The log_softmax lp and the one-hot rows of the targets are saved
1662# for the VJP. Logits of another rank are refused with NotLogits (107),
1663# targets that do not fit with TargetOutOfRange (108), an unknown id with
1664# UnknownNode; the tape comes back unchanged.
1665def cross_entropy(tape: Tape, +a: Nat, +targets: List<&2, U32>) -> Tape & Result<&2, &2, Error, Nat>:
1666  Tape{entries} = tape
1667  entropy_found(a, targets, value_counted(a, counted(entries)))