def slice_axis_plan

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

import Base
import bendtorch-core@0.1.0.0/src/shape/core.bend as Shape
import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed
import bendtorch-core@0.1.0.0/src/tensor/slice.bend as Slice

1455def slice_axis_plan(+shape: Shape.Shape, +axis: U32, +slicing: Slicing) -> Result<&2, &2, Slice.Error, Slice.Planned>:
1456  match slicing:
1457    case Sliced{dim, start, stop, +step}:
1458      Slice.plan(shape, axis, Signed.bound(slice_extent(shape, axis), start), Signed.bound(slice_extent(shape, axis), stop), step)
1459    case Narrowed{dim, start, +length}:
1460      +first = Signed.wrap(slice_extent(shape, axis), start)
1461      Slice.plan(shape, axis, first, U32.add(first, length), 1)
1462    case Selected{dim, index}:
1463      +first = Signed.wrap(slice_extent(shape, axis), index)
1464      Slice.plan(shape, axis, first, U32.add(first, 1), 1)