Tensor.sum_axis_keep: def of bendtorch@1.3.0.0
Import: import bendtorch@1.3.0.0/tensor.bend as Tensor
Lines 407-411 of bendtorch@1.3.0.0/tensor.bend: https://jend.leeeet.dev/read.txt?path=bendtorch@1.3.0.0/tensor.bend&start_line=407&end_line=411

import Base

# the sum over axis k with the axis kept (size 1): the same numbers as sum_axis, in the same layout
def sum_axis_keep(+s: List<&2, Nat>, +k: Nat, x: Tensor<s>) -> TP<s, setd(s, k, 1n)>:
  match x:
    case Tensor{d}:
      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)))

Uses (13):
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Array.get
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=F32
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=List
  https://jend.leeeet.dev/definition.txt?path=Base/base.bend&name=Nat
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=TP
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=Tensor
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=drop_map
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=new
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=numel
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=remove
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=scatter_l
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=setd
  https://jend.leeeet.dev/definition.txt?path=bendtorch@1.3.0.0/tensor.bend&name=tp_of