def sum_axis_keep

bendtorch@1.3.0.0 · tensor.bend:407

import Base

407# the sum over axis k with the axis kept (size 1): the same numbers as sum_axis, in the same layout
408def sum_axis_keep(+s: List<&2, Nat>, +k: Nat, x: Tensor<s>) -> TP<s, setd(s, k, 1n)>:
409  match x:
410    case Tensor{d}:
411      tp_of(s, setd(s, k, 1n), scatter_l(numel(s), 0, drop_map(s, k), Array.get(F32, d, 0), Array.get(F32, new(numel(remove(s, k))), 0)))