def select

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

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

320# PyTorch's input.select(dim, index) at a signed axis and a signed index,
321# wrapped among the extent: an index outside -extent..extent-1 stays at or
322# past the extent, so the unsigned select refuses it with OutOfRange after
323# its axis check.
324def select(input: Core.Dynamic, dim: Position, index: Position) -> Core.Dynamic & Result<&2, &1, Slice.Error, Core.Dynamic>:
325  Core.Packed{+shape, tensor} = input
326  +axis = wrap(Shape.rank(shape), dim)
327  Slice.select(Core.Packed{shape, tensor}, axis, wrap(extent(Shape.dimensions(shape), axis), index))