Trait IndexingOps 

pub trait IndexingOps<R>
where
    R: Runtime,
{
    // Provided methods
    fn argmax(
        &self,
        a: &Tensor<R>,
        dim: usize,
        keepdim: bool,
    ) -> Result<Tensor<R>, Error> { ... }
    fn argmin(
        &self,
        a: &Tensor<R>,
        dim: usize,
        keepdim: bool,
    ) -> Result<Tensor<R>, Error> { ... }
    fn gather(
        &self,
        a: &Tensor<R>,
        dim: usize,
        index: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn scatter(
        &self,
        a: &Tensor<R>,
        dim: usize,
        index: &Tensor<R>,
        src: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn index_select(
        &self,
        a: &Tensor<R>,
        dim: usize,
        index: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn index_put(
        &self,
        a: &Tensor<R>,
        dim: usize,
        index: &Tensor<R>,
        src: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn take(
        &self,
        tensor: &Tensor<R>,
        indices: &Tensor<R>,
    ) -> Result<Tensor<R>, Error>
       where R: Runtime<DType = DType> { ... }
    fn put(
        &self,
        tensor: &Tensor<R>,
        indices: &Tensor<R>,
        values: &Tensor<R>,
    ) -> Result<Tensor<R>, Error>
       where R: Runtime<DType = DType> { ... }
    fn masked_select(
        &self,
        a: &Tensor<R>,
        mask: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn masked_fill(
        &self,
        a: &Tensor<R>,
        mask: &Tensor<R>,
        value: f64,
    ) -> Result<Tensor<R>, Error> { ... }
    fn embedding_lookup(
        &self,
        embeddings: &Tensor<R>,
        indices: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn scatter_reduce(
        &self,
        dst: &Tensor<R>,
        dim: usize,
        index: &Tensor<R>,
        src: &Tensor<R>,
        op: ScatterReduceOp,
        include_self: bool,
    ) -> Result<Tensor<R>, Error> { ... }
    fn gather_nd(
        &self,
        input: &Tensor<R>,
        indices: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn bincount(
        &self,
        input: &Tensor<R>,
        weights: Option<&Tensor<R>>,
        minlength: usize,
    ) -> Result<Tensor<R>, Error> { ... }
    fn gather_2d(
        &self,
        input: &Tensor<R>,
        rows: &Tensor<R>,
        cols: &Tensor<R>,
    ) -> Result<Tensor<R>, Error> { ... }
    fn slice_assign(
        &self,
        dst: &Tensor<R>,
        src: &Tensor<R>,
        dim: usize,
        start: usize,
    ) -> Result<Tensor<R>, Error> { ... }
}

Indexing operations

Provided Methods

fn argmax( &self, a: &Tensor<R>, dim: usize, keepdim: bool, ) -> Result<Tensor<R>, Error>

Argmax: returns indices of maximum values along a dimension.

Returns a tensor of I64 indices indicating the position of the maximum value along the specified dimension. The output shape is the input shape with the specified dimension removed (or kept as size 1 if keepdim=true).

Arguments

• a - Input tensor
• dim - Dimension along which to find the maximum index
• keepdim - If true, the reduced dimension is retained with size 1

Returns

Tensor of I64 containing indices of maximum values

fn argmin( &self, a: &Tensor<R>, dim: usize, keepdim: bool, ) -> Result<Tensor<R>, Error>

Argmin: returns indices of minimum values along a dimension.

Returns a tensor of I64 indices indicating the position of the minimum value along the specified dimension. The output shape is the input shape with the specified dimension removed (or kept as size 1 if keepdim=true).

Arguments

• a - Input tensor
• dim - Dimension along which to find the minimum index
• keepdim - If true, the reduced dimension is retained with size 1

Returns

Tensor of I64 containing indices of minimum values

fn gather( &self, a: &Tensor<R>, dim: usize, index: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Gather elements along a dimension using an index tensor.

For a 3D tensor with dim=1: `out[i][j][k] = input[i][index[i][j][k]][k]`

Arguments

• a - Input tensor
• dim - Dimension along which to gather
• index - Index tensor (I32 or I64) with same number of dimensions as input

Returns

Tensor with same shape as index tensor, same dtype as input

fn scatter( &self, a: &Tensor<R>, dim: usize, index: &Tensor<R>, src: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Scatter values into a tensor at positions specified by an index tensor.

Creates a new tensor (copy of a) with values from src scattered at positions specified by index along dimension dim.

For a 3D tensor with dim=1: `out[i][index[i][j][k]][k] = src[i][j][k]`

Arguments

• a - Input tensor (values to scatter into)
• dim - Dimension along which to scatter
• index - Index tensor (I32 or I64) specifying scatter positions
• src - Source tensor with values to scatter

Returns

New tensor with scattered values

fn index_select( &self, a: &Tensor<R>, dim: usize, index: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Select elements along a dimension using a 1D index tensor.

Simpler than gather - the index tensor is 1D and applies to all positions in the specified dimension.

Arguments

• a - Input tensor
• dim - Dimension along which to select
• index - 1D index tensor (I32 or I64) of length m

Returns

Tensor with dimension dim having size m (length of index)

Bounds Checking

Returns IndexOutOfBounds error if any index is negative or >= dim_size. Indices must be in the range [0, dim_size). Negative indices are not supported.

fn index_put( &self, a: &Tensor<R>, dim: usize, index: &Tensor<R>, src: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Put values at specified indices along a dimension.

This is the inverse of index_select - it assigns values from src into positions specified by index along dimension dim.

Example

# Replace row 2 of a [5, 3] matrix with new values:
a = [[1, 2, 3],
     [4, 5, 6],
     [7, 8, 9],      # <- row 2 will be replaced
     [10, 11, 12],
     [13, 14, 15]]
index = [2]          # indices along dim 0
src = [[100, 200, 300]]

result = index_put(a, 0, index, src)
# result = [[1, 2, 3],
#           [4, 5, 6],
#           [100, 200, 300],  # replaced!
#           [10, 11, 12],
#           [13, 14, 15]]

Arguments

• a - Input tensor to modify (copied, not mutated)
• dim - Dimension along which to put values
• index - 1D index tensor (I32 or I64) specifying positions
• src - Source tensor with values to insert. Shape must match a except at dim where it must equal index.numel()

Returns

New tensor with values at indexed positions replaced by src

Bounds Checking

Returns IndexOutOfBounds error if any index is negative or >= dim_size. Indices must be in the range [0, dim_size). Negative indices are not supported.

fn take( &self, tensor: &Tensor<R>, indices: &Tensor<R>, ) -> Result<Tensor<R>, Error>

where R: Runtime<DType = DType>,

Take values from a tensor using flat indices.

The input tensor is treated as flattened 1D storage, and values are gathered at positions specified by indices. The output shape matches indices.shape().

Arguments

• tensor - Input tensor to gather from
• indices - Index tensor (I32 or I64) containing flat indices

Returns

Tensor of shape indices.shape() with gathered values

fn put( &self, tensor: &Tensor<R>, indices: &Tensor<R>, values: &Tensor<R>, ) -> Result<Tensor<R>, Error>

where R: Runtime<DType = DType>,

Put values into a tensor at flat indices (functional, non-mutating).

Returns a new tensor with values written at positions specified by indices, treating the input tensor as flattened 1D storage.

Arguments

• tensor - Input tensor to update
• indices - Index tensor (I32 or I64) containing flat indices
• values - Values to write. Must have the same number of elements as indices.

Returns

New tensor with the same shape as tensor and updated values

fn masked_select( &self, a: &Tensor<R>, mask: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Select elements where mask is true, returning a flattened 1D tensor.

Arguments

• a - Input tensor
• mask - Boolean mask tensor (U8: 0=false, non-zero=true), must be broadcastable to a

Returns

1D tensor containing only elements where mask is true

fn masked_fill( &self, a: &Tensor<R>, mask: &Tensor<R>, value: f64, ) -> Result<Tensor<R>, Error>

Fill elements where mask is true with a scalar value.

Arguments

• a - Input tensor
• mask - Boolean mask tensor (U8: 0=false, non-zero=true), must be broadcastable to a
• value - Value to fill where mask is true

Returns

New tensor with masked positions filled with value

fn embedding_lookup( &self, embeddings: &Tensor<R>, indices: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Look up embeddings from an embedding table using indices.

This is the standard embedding lookup operation used in neural networks for word embeddings, entity embeddings, etc. It is equivalent to index_select(embeddings, 0, indices) but optimized for the common case where the embedding table is 2D and indices index into the first dimension.

Algorithm

For each index value i in the indices tensor:

output[..., i, :] = embeddings[indices[..., i], :]

The output shape is indices.shape() + [embedding_dim] where embedding_dim is embeddings.shape()[1].

Arguments

• embeddings - 2D embedding table of shape `[vocab_size, embedding_dim]`
• indices - Index tensor of any shape containing indices into the embedding table. Must be I32 or I64. Values must be in range `[0, vocab_size)`.

Returns

Tensor of shape `indices.shape() + [embedding_dim]` containing the looked-up embeddings.

Example

embeddings = [[1.0, 2.0],   # word 0
              [3.0, 4.0],   # word 1
              [5.0, 6.0]]   # word 2
indices = [2, 0, 1]

output = [[5.0, 6.0],   # word 2
          [1.0, 2.0],   # word 0
          [3.0, 4.0]]   # word 1

Errors

• ShapeMismatch - if embeddings is not 2D
• DTypeMismatch - if indices is not I32 or I64
• Index out of bounds results in undefined behavior (implementation may return zeros)

Performance

On GPU, this operation is memory-bound and optimized for coalesced reads from the embedding table. Each thread handles one index lookup and writes a full embedding vector.

fn scatter_reduce( &self, dst: &Tensor<R>, dim: usize, index: &Tensor<R>, src: &Tensor<R>, op: ScatterReduceOp, include_self: bool, ) -> Result<Tensor<R>, Error>

Scatter values with reduction into a destination tensor.

Unlike regular scatter which overwrites values, scatter_reduce applies a reduction operation when multiple source values scatter to the same destination index.

Algorithm

For each position in src:

dst[..., index[...], ...] = reduce(dst[..., index[...], ...], src[...])

Where reduce is determined by the op parameter.

Arguments

• dst - Destination tensor to scatter into (used as initial values)
• dim - Dimension along which to scatter
• index - Index tensor (I32 or I64) specifying scatter positions
• src - Source tensor with values to scatter
• op - Reduction operation to apply (Sum, Mean, Max, Min, Prod)
• include_self - If true, include dst values in reduction; if false, initialize destination positions from src only

Returns

New tensor with scattered and reduced values

Example

dst = [0, 0, 0, 0]
index = [0, 0, 2]
src = [1, 2, 3]
scatter_reduce(dst, 0, index, src, Sum, include_self=true)
# Result: [3, 0, 3, 0]  // src[0]+src[1]=3 at index 0, src[2]=3 at index 2

fn gather_nd( &self, input: &Tensor<R>, indices: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Gather elements using N-dimensional indices.

Unlike regular gather which gathers along a single dimension, gather_nd uses an index tensor where the last dimension specifies coordinates into the input tensor.

Algorithm

If indices has shape `[..., M]` and input has N dimensions:

• If M == N: output has shape `indices.shape()[:-1]`
• If M < N: output has shape `indices.shape()[:-1] + input.shape()[M:]`

Each index vector `indices[..., :]` specifies coordinates for the first M dimensions of input.

Arguments

• input - Input tensor to gather from
• indices - Index tensor where last dimension contains coordinates

Returns

Tensor with gathered values

Example

input = [[0, 1], [2, 3]]  # shape [2, 2]
indices = [[0, 0], [1, 1]]  # shape [2, 2], last dim=2 means full coordinates
gather_nd(input, indices)
# Result: [0, 3]  # input[0,0]=0, input[1,1]=3

indices = [[0], [1]]  # shape [2, 1], last dim=1 means gather rows
gather_nd(input, indices)
# Result: [[0, 1], [2, 3]]  # input[0,:], input[1,:]

fn bincount( &self, input: &Tensor<R>, weights: Option<&Tensor<R>>, minlength: usize, ) -> Result<Tensor<R>, Error>

Count occurrences of each value in an integer tensor.

Returns a histogram where `output[i]` contains the count (or weighted sum) of how many times value i appears in the input.

Arguments

• input - 1D integer tensor with non-negative values (I32 or I64)
• weights - Optional weights tensor, same shape as input. If provided, the output is the sum of weights for each bin instead of counts.
• minlength - Minimum length of the output tensor. Useful when the maximum value is known ahead of time.

Returns

1D tensor of length `max(max(input)+1, minlength)` containing counts or weighted sums.

Example

input = [0, 1, 1, 3, 2, 1, 3]
bincount(input, None, 0)
# Result: [1, 3, 1, 2]  // counts: 0->1, 1->3, 2->1, 3->2

weights = [0.5, 1.0, 1.5, 2.0, 1.0, 0.5, 3.0]
bincount(input, Some(weights), 0)
# Result: [0.5, 3.0, 1.0, 5.0]  // weighted sums per bin

Errors

• ShapeMismatch - if input is not 1D or weights shape doesn’t match input
• DTypeMismatch - if input is not an integer type
• InvalidValue - if input contains negative values

fn gather_2d( &self, input: &Tensor<R>, rows: &Tensor<R>, cols: &Tensor<R>, ) -> Result<Tensor<R>, Error>

Gather elements from a 2D matrix using row and column index vectors.

For each index i, extracts `input[rows[i], cols[i]]`.

This is a specialized gather operation optimized for sparse matrix applications where you need to extract values at specific (row, col) coordinates.

Algorithm

output[i] = input[rows[i], cols[i]]  for i in 0..len(rows)

Arguments

• input - 2D input tensor of shape `[nrows, ncols]`
• rows - 1D index tensor (I32 or I64) specifying row indices
• cols - 1D index tensor (I32 or I64) specifying column indices

Returns

1D tensor of length `rows.numel()` with gathered values. Same dtype as input.

Example

input = [[1, 2, 3],
         [4, 5, 6],
         [7, 8, 9]]
rows = [0, 1, 2, 0]
cols = [0, 1, 2, 2]

gather_2d(input, rows, cols)
# Result: [1, 5, 9, 3]  // input[0,0], input[1,1], input[2,2], input[0,2]

Errors

• ShapeMismatch - if input is not 2D or rows/cols have different lengths
• DTypeMismatch - if rows or cols are not I32 or I64
• IndexOutOfBounds - if any (row, col) pair is out of bounds

fn slice_assign( &self, dst: &Tensor<R>, src: &Tensor<R>, dim: usize, start: usize, ) -> Result<Tensor<R>, Error>

Assign src into a slice of dst along dimension dim starting at start.

Returns a new tensor equal to dst except that the region dst[..., start..start+src.shape[dim], ...] is replaced by src.

Arguments

• dst - Destination tensor
• src - Source tensor. Must have same shape as dst except at dim, where src.shape[dim] + start <= dst.shape[dim]
• dim - Dimension along which to assign
• start - Starting index in dst along dim

Returns

New tensor with the slice replaced

Implementors

impl IndexingOps<CpuRuntime> for CpuClient

IndexingOps implementation for CPU runtime.