def gather_axis

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

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/layout.bend as Layout

181def gather_axis(+dimensions: List<&2, U32>, +axis: Nat, +index: List<&2, U32>, +entries: List<&2, U32>, in_rank: Bool) ->
182  Result<&2, &2, Error, Planned>:
183  match in_rank:
184    case False{}:
185      Fail{LayoutError{Layout.AxisOutOfBounds{}}}
186    case True{}:
187      gather_ranked(dimensions, axis, index, entries, Nat.is_eq(Shape.rank_dimensions(index), Shape.rank_dimensions(dimensions)))

Uses (11)

Used by (6)