Trait TensorOps 

pub trait TensorOps: Send + Sync {
    // Required methods
    fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
    fn add(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
    fn mul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
    fn attention(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        mask: Option<&Tensor>,
        scale: Option<f32>,
    ) -> Result<Tensor>;
    fn layer_norm(
        &self,
        input: &Tensor,
        weight: &Tensor,
        bias: Option<&Tensor>,
        eps: f32,
    ) -> Result<Tensor>;
    fn gelu(&self, input: &Tensor) -> Result<Tensor>;
    fn silu(&self, input: &Tensor) -> Result<Tensor>;
    fn softmax(&self, input: &Tensor, dim: isize) -> Result<Tensor>;
    fn embedding(&self, indices: &Tensor, weight: &Tensor) -> Result<Tensor>;
    fn zeros(
        &self,
        shape: &[usize],
        dtype: DataType,
        device: &Device,
    ) -> Result<Tensor>;
    fn randn(
        &self,
        shape: &[usize],
        dtype: DataType,
        device: &Device,
    ) -> Result<Tensor>;
    fn exp(&self, input: &Tensor) -> Result<Tensor>;
    fn normalize(&self, input: &Tensor, p: i32, dim: i32) -> Result<Tensor>;
    fn concat(&self, tensors: &[&Tensor], dim: usize) -> Result<Tensor>;
    fn rms_norm(
        &self,
        input: &Tensor,
        weight: &Tensor,
        eps: f32,
    ) -> Result<Tensor>;
    fn scale(&self, input: &Tensor, factor: f32) -> Result<Tensor>;
    fn sigmoid(&self, input: &Tensor) -> Result<Tensor>;
    fn topk(
        &self,
        input: &Tensor,
        k: usize,
        dim: i64,
    ) -> Result<(Tensor, Tensor)>;
    fn conv1d(
        &self,
        input: &Tensor,
        weight: &Tensor,
        bias: Option<&Tensor>,
        stride: usize,
        padding: usize,
    ) -> Result<Tensor>;
    fn tanh(&self, input: &Tensor) -> Result<Tensor>;
    fn sub(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
    fn clamp(&self, input: &Tensor, min: f32, max: f32) -> Result<Tensor>;
    fn gather(
        &self,
        input: &Tensor,
        dim: usize,
        indices: &Tensor,
    ) -> Result<Tensor>;
    fn scatter(
        &self,
        input: &Tensor,
        dim: usize,
        indices: &Tensor,
        src: &Tensor,
    ) -> Result<Tensor>;
    fn flash_attention(
        &self,
        query: &Tensor,
        key: &Tensor,
        value: &Tensor,
        scale: f32,
        causal: bool,
    ) -> Result<Tensor>;
    fn fused_swiglu(&self, gate: &Tensor, up: &Tensor) -> Result<Tensor>;
    fn fused_residual_rms_norm(
        &self,
        residual: &Tensor,
        hidden: &Tensor,
        weight: &Tensor,
        eps: f32,
    ) -> Result<Tensor>;
}

Core tensor operations trait - device-agnostic interface

Required Methods

fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>

Matrix multiplication: C = A @ B

fn add(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>

Element-wise addition: C = A + B

fn mul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>

Element-wise multiplication: C = A * B

fn attention( &self, query: &Tensor, key: &Tensor, value: &Tensor, mask: Option<&Tensor>, scale: Option<f32>, ) -> Result<Tensor>

Scaled dot-product attention

fn layer_norm( &self, input: &Tensor, weight: &Tensor, bias: Option<&Tensor>, eps: f32, ) -> Result<Tensor>

Layer normalization

fn gelu(&self, input: &Tensor) -> Result<Tensor>

GELU activation

fn silu(&self, input: &Tensor) -> Result<Tensor>

SiLU/Swish activation

fn softmax(&self, input: &Tensor, dim: isize) -> Result<Tensor>

Softmax

fn embedding(&self, indices: &Tensor, weight: &Tensor) -> Result<Tensor>

Embedding lookup

fn zeros( &self, shape: &[usize], dtype: DataType, device: &Device, ) -> Result<Tensor>

Create tensor with zeros

fn randn( &self, shape: &[usize], dtype: DataType, device: &Device, ) -> Result<Tensor>

Create tensor with random values

fn exp(&self, input: &Tensor) -> Result<Tensor>

Exponential function

fn normalize(&self, input: &Tensor, p: i32, dim: i32) -> Result<Tensor>

L2 normalization

fn concat(&self, tensors: &[&Tensor], dim: usize) -> Result<Tensor>

Concatenate tensors along a dimension

fn rms_norm(&self, input: &Tensor, weight: &Tensor, eps: f32) -> Result<Tensor>

RMS normalization

fn scale(&self, input: &Tensor, factor: f32) -> Result<Tensor>

Scale tensor by a factor

fn sigmoid(&self, input: &Tensor) -> Result<Tensor>

Sigmoid activation: 1 / (1 + exp(-x))

fn topk(&self, input: &Tensor, k: usize, dim: i64) -> Result<(Tensor, Tensor)>

Top-k operation: returns (values, indices) for top k elements along dimension

fn conv1d( &self, input: &Tensor, weight: &Tensor, bias: Option<&Tensor>, stride: usize, padding: usize, ) -> Result<Tensor>

1D convolution

fn tanh(&self, input: &Tensor) -> Result<Tensor>

Tanh activation

fn sub(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>

Element-wise subtraction

fn clamp(&self, input: &Tensor, min: f32, max: f32) -> Result<Tensor>

Clamp values to a range

fn gather(&self, input: &Tensor, dim: usize, indices: &Tensor) -> Result<Tensor>

Gather elements along dimension using indices

fn scatter( &self, input: &Tensor, dim: usize, indices: &Tensor, src: &Tensor, ) -> Result<Tensor>

Scatter elements along dimension using indices

fn flash_attention( &self, query: &Tensor, key: &Tensor, value: &Tensor, scale: f32, causal: bool, ) -> Result<Tensor>

Fused scaled dot-product attention (Flash Attention pattern) Computes: softmax(Q @ K^T / sqrt(d_k)) @ V in a memory-efficient manner Works for all attention-based models: LLaMA, Qwen, Gemma, Mistral, etc.

fn fused_swiglu(&self, gate: &Tensor, up: &Tensor) -> Result<Tensor>

Fused SwiGLU activation: silu(gate) * up Used by LLaMA, Qwen, Mistral, etc.

fn fused_residual_rms_norm( &self, residual: &Tensor, hidden: &Tensor, weight: &Tensor, eps: f32, ) -> Result<Tensor>

Fused residual add + RMS norm Computes: rms_norm(residual + hidden, weight, eps)

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors

impl TensorOps for CpuTensorOpsImpl

impl TensorOps for GpuTensorOpsImpl

Basic GPU implementation of tensor operations For now, this just delegates to CPU implementation

impl TensorOps for TensorDispatcher