def sum_dims

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

import Base
import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
import bendtorch-core@0.1.0.0/src/autograd/sum_dims.bend as SumDims

1412# Records PyTorch's a.sum(dim=axes, keepdim=keepdim), SumDims.sum_dims of a
1413# copy of a's value: the kept shape when keepdim is True, else the dropped
1414# one. The axes are PyTorch's signed ones: Last{0} is -1, a's last axis. An
1415# empty axis list sums every axis. An axis outside -limit..limit-1, where
1416# limit is max(rank, 1), so a rank-0 a takes -1 and 0, is refused with
1417# AxisOutOfBounds (13), two axes that name one axis with DuplicateAxis (14),
1418# an unknown id with UnknownNode; the tape comes back unchanged.
1419def sum_dims(tape: Tape, +a: Nat, +axes: List<&2, Signed.Position>, keepdim: Bool) -> Tape & Result<&2, &2, Error, Nat>:
1420  Tape{entries} = tape
1421  dims_found(a, axes, keepdim, value_counted(a, counted(entries)))