~/jendhub

def sum_axis_bwd

bendtorch@1.3.0.0 · tensor.bend:283

import bendtorch@1.3.0.0/tensor.bend as Tensor
import Base

283# the gradient of sum_axis: dy (shape remove(s, k)) copied along axis k; dy comes back
284def sum_axis_bwd(+s: List<&2, Nat>, +k: Nat, dy: Tensor<remove(s, k)>) -> TP<remove(s, k), s>:
285  match dy:
286    case Tensor{d}:
287      tp_of(remove(s, k), s, gather_l(numel(s), 0, drop_map(s, k), Array.get(F32, d, imap(drop_map(s, k), 0)), new(numel(s))))

Uses (13)

Used by (1)