def gather_plan

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

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

197# The plan of gather(dim, index), or its refusal, in this order: an entry
198# count other than the index shape's element count, the axis, the index's
199# rank, its extents at the other axes, then its entries. The element count
200# is the Shape's own, so no product of its extents is formed here.
201def gather_plan(+input: Shape.Shape, dim: Signed.Position, +index: Index) -> Result<&2, &2, Error, Planned>:
202  Index{+shape, +entries} = index
203  +axis = U32.to_nat(Signed.wrap(Shape.rank(input), dim))
204  gather_counted(Shape.dimensions(input), axis, Shape.dimensions(shape), entries,
205    Nat.is_eq(count(entries), Shape.elements(shape)))