Trait TensorIndexingPrims

Source

pub trait TensorIndexingPrims<Alg: Algebra>: Sized {
    type Plan;
    type Context;

    // Required methods
    fn plan(
        ctx: &mut Self::Context,
        desc: &IndexingPrimsDescriptor,
        shapes: &[&[usize]],
    ) -> Result<Self::Plan>;
    fn execute(
        ctx: &mut Self::Context,
        plan: &Self::Plan,
        inputs: &[&Tensor<Alg::Scalar>],
        indices: &Tensor<i64>,
        output: &mut Tensor<Alg::Scalar>,
    ) -> Result<()>;
    fn has_indexing_support(desc: IndexingPrimsDescriptor) -> bool;
}

Expand description

Indexing execution protocol family.

Provides index-based selection, gathering, and scattering of tensor elements. Unlike the scalar and analytic families, indexing does not use alpha/beta scaling — operations are pure data movement.

The indices tensor uses i64 elements representing positions along the target axis. Negative indices are not supported and will produce errors.

§Examples

use tenferro_algebra::Standard;
use tenferro_prims::{CpuBackend, CpuContext, IndexingPrimsDescriptor, TensorIndexingPrims};

let mut ctx = CpuContext::new(1);
let desc = IndexingPrimsDescriptor::IndexSelect { axis: 0 };
let _plan = <CpuBackend as TensorIndexingPrims<Standard<f64>>>::plan(
    &mut ctx,
    &desc,
    &[&[3, 2], &[2], &[2, 2]],
)
.unwrap();

Required Associated Types§

Source

type Plan

Backend-specific execution plan.

Source

type Context

Backend execution context.

Required Methods§

Source

fn plan( ctx: &mut Self::Context, desc: &IndexingPrimsDescriptor, shapes: &[&[usize]], ) -> Result<Self::Plan>

Plan an indexing operation for the given input/index/output shapes.

shapes contains: [input_shape, index_shape, output_shape].

Source

fn execute( ctx: &mut Self::Context, plan: &Self::Plan, inputs: &[&Tensor<Alg::Scalar>], indices: &Tensor<i64>, output: &mut Tensor<Alg::Scalar>, ) -> Result<()>

Execute a previously planned indexing operation.

For IndexSelect and Gather: reads from inputs[0] using indices, writes into output.
For Scatter: reads from inputs[0] (source values) using indices, writes into output.
For IndexPut: reads from inputs[0] (values to put) using indices, writes into output.