def select_plan

bendtorch-core@0.1.0.0 · src/tensor/gather.bend:159

import bendtorch-core@0.1.0.0/src/tensor/gather.bend as Gather
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

159# The plan of index_select(dim, index), or its refusal.
160def select_plan(+input: Shape.Shape, dim: Signed.Position, +entries: List<&2, U32>) -> Result<&2, &2, Error, Planned>:
161  +dimensions = Shape.dimensions(input)
162  +axis = U32.to_nat(Signed.wrap(Shape.rank(input), dim))
163  selecting(dimensions, axis, entries, Nat.is_lt(axis, Shape.rank_dimensions(dimensions)))

Uses (14)

Used by (6)