def narrow

bendtorch-core@0.1.0.0 · src/tensor/signed.bend:346

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/slice.bend as Slice

346# PyTorch's input.narrow(dim, start, length) at a signed axis with a signed
347# start. PyTorch takes a start in -extent..extent: Last{k} with k below the
348# extent is extent - 1 - k, and one before -extent is OutOfRange. A First
349# start past the extent is refused by the unsigned narrow, which checks
350# start + length against the extent, with the same code. A length is a
351# U32, so PyTorch's negative lengths cannot be written.
352def narrow(input: Core.Dynamic, dim: Position, +start: Position, +length: U32) ->
353  Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
354  Core.Packed{+shape, tensor} = input
355  +axis = wrap(Shape.rank(shape), dim)
356  +n = extent(Shape.dimensions(shape), axis)
357  narrow_started(shape, tensor, axis, wrap(n, start), length, before(n, start))