Struct TensorDispatcher 

pub struct TensorDispatcher { /* private fields */ }

Tensor operation dispatcher - automatically selects CPU/GPU implementation

Implementations

impl TensorDispatcher

pub fn new() -> Self

Trait Implementations

impl TensorOps for TensorDispatcher

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)