def index_select

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

import Base
import bendtorch-core@0.1.0.0/src/tensor/core.bend as Core
import bendtorch-core@0.1.0.0/src/tensor/signed.bend as Signed

436# PyTorch's input.index_select(dim, index). The input owner always comes back
437# unchanged next to the result.
438def index_select(input: Core.Dynamic, dim: Signed.Position, +entries: List<&2, U32>) -> Core.Dynamic & Result<&2, &1, Error, Core.Dynamic>:
439  Core.Packed{+shape, tensor} = input
440  copied(shape, tensor, entries, select_plan(shape, dim, entries))