~/jendhub

def sum_axis

bendtorch@1.3.0.0 · tensor.bend:277

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

277# x summed over axis k: shape remove(s, k); x comes back
278def sum_axis(+s: List<&2, Nat>, +k: Nat, x: Tensor<s>) -> TP<s, remove(s, k)>:
279  match x:
280    case Tensor{d}:
281      tp_of(s, remove(s, k), scatter_l(numel(s), 0, drop_map(s, k), Array.get(F32, d, 0), Array.get(F32, new(numel(remove(s, k))), 0)))

Uses (12)

Used by (1)