~/jendhub

def vjp_sum_dims

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

import bendtorch-core@0.1.0.0/src/autograd/tensor_backward.bend as TensorBackward
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

307# d(sum of a over some axes)/da: g read with the kept shape, the reduced axes
308# put back as 1 (PyTorch's unsqueeze when keepdim is False), then broadcast
309# to a's shape, as the Sum VJP does.
310def vjp_sum_dims(+a: Nat, +input: Shape.Shape, +kept: Shape.Shape, g: Core.Dynamic) -> Pass:
311  Core.Packed{+shape, tensor} = g
312  unsqueezed(a, input, kept, shape, tensor, Shape.bool_case(Nat.is_eq(Shape.elements(shape), Shape.elements(kept))))

Uses (8)

Used by (1)