def dims_shaped

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

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
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

1397# The signed axes, each wrapped among the operand's axes (its rank, or 1 at
1398# rank 0), as PyTorch wraps them, before the plan checks them.
1399def dims_shaped(+a: Nat, +axes: List<&2, Signed.Position>, keepdim: Bool, entries: List<&1, Entry>, pair: Core.Dynamic & Shape.Shape) ->
1400  Tape & Result<&2, &2, Error, Nat>:
1401  (input, +shape) = pair
1402  dims_planned(a, entries, input, shape, keepdim, SumDims.plan(shape, Signed.wrap_all(SumDims.limit(Shape.dimensions(shape)), axes)))