Struct Tensor 

pub struct Tensor {
    pub shape: Vec<usize>,
    pub strides: Vec<usize>,
    /* private fields */
}

A general N-D f32 tensor: an Arc-shared device buffer viewed through (shape, strides, offset).

Fields

shape: Vec<usize>

strides: Vec<usize>

Implementations

impl Tensor

pub fn to_half(&self, dtype: DType) -> Half

Pack this f32 tensor down to half precision (round-to-nearest-even), on-device.

impl Tensor

pub async fn quantize_i8(&self) -> QTensor

Symmetric per-tensor int8 quantization (scale = max|x|/127). Async: the scalar scale is read back so the quantized matmul can fold both scales into one small buffer (WebGPU allows only 4 storage buffers per shader — scalars ride in the info buffer instead of their own bindings).

impl Tensor

pub fn quantize_rowwise(&self, bits: u32) -> QRow

Per-row symmetric quantization of a 2D matrix at 4 or 8 bits (scale = max|row|/(2^(bits-1)−1)).

impl Tensor

pub fn matmul_qweight(&self, w: &QRow) -> Tensor

Weight-only quantized matmul (the efficient-inference path): x [rows, in] · Wᵀ where W is a per-row-quantized [out, in] matrix that stays packed in memory — dequantized on the fly in the kernel. Returns [rows, out]. This is W4A16/W8A16-style: activations f32, weights int4/int8.

impl Tensor

pub fn quantize_ternary(&self) -> Ternary

Quantize a 2D [out,in] weight to ternary {−1,0,+1} with per-row absmean scale (BitNet-style).

pub fn matmul_ternary(&self, w: &Ternary) -> Tensor

Multiply-free ternary matmul: x [rows,in] · Wᵀ where W is ternary [out,in]. Returns [rows,out].

impl Tensor

pub fn numel(&self) -> usize

pub fn rank(&self) -> usize

pub fn is_contiguous(&self) -> bool

pub fn from_vec(ctx: &Arc<Context>, data: &[f32], shape: &[usize]) -> Tensor

pub fn zeros(ctx: &Arc<Context>, shape: &[usize]) -> Tensor

pub async fn to_vec(&self) -> Vec<f32>

Materialize (contiguous) and read back to host in logical row-major order.

pub fn reshape(&self, shape: &[usize]) -> Tensor

pub fn permute(&self, perm: &[usize]) -> Tensor

pub fn transpose(&self, a: usize, b: usize) -> Tensor

pub fn broadcast_to(&self, shape: &[usize]) -> Tensor

Broadcast to a larger shape (right-aligned); broadcast dims get stride 0.

pub fn contiguous(&self) -> Tensor

Materialize a (possibly strided/broadcast) view into a fresh contiguous buffer.

pub fn add(&self, o: &Tensor) -> Tensor

pub fn sub(&self, o: &Tensor) -> Tensor

pub fn mul(&self, o: &Tensor) -> Tensor

pub fn div(&self, o: &Tensor) -> Tensor

pub fn maximum(&self, o: &Tensor) -> Tensor

pub fn exp(&self) -> Tensor

pub fn neg(&self) -> Tensor

pub fn relu(&self) -> Tensor

pub fn sqrt(&self) -> Tensor

pub fn relu_mask(&self) -> Tensor

pub fn abs(&self) -> Tensor

pub fn sigmoid(&self) -> Tensor

pub fn silu(&self) -> Tensor

pub fn gelu(&self) -> Tensor

pub fn log(&self) -> Tensor

pub fn relu2(&self) -> Tensor

pub fn scalar(&self, s: f32) -> Tensor

pub fn softmax(&self, axis: usize) -> Tensor

Softmax over axis (fused: per-row max/exp/sum/div in one kernel).

pub fn rmsnorm(&self, weight: &Tensor, eps: f32) -> Tensor

RMSNorm over the last dim: x/sqrt(mean(x²)+eps)·weight (fused).

pub fn rope( &self, n_heads: usize, head_dim: usize, base: f32, offset: usize, ) -> Tensor

Rotary position embedding (NeoX rotate-half) on a [T, n_heads·head_dim] tensor.

pub fn rope_3d( &self, n_heads: usize, head_dim: usize, base: f32, gt: usize, gh: usize, gw: usize, ) -> Tensor

3D rotary position embedding (V-JEPA 2): head_dim split into 3 groups (temporal/height/width), each rotated by the token’s coordinate along that axis. self is [T, n_heads·head_dim], T=gt·gh·gw.

pub fn depthwise_conv1d_causal(&self, weight: &Tensor, l: usize) -> Tensor

Causal depthwise conv1d — the LFM2 / Liquid AI short-conv mixer. self is [T, C] (sequence × channels), weight is [C, L] (per-channel kernel of length L). Causal: out[t] sees only t-L+1..t.

pub fn matmul_bt(&self, w: &Tensor) -> Tensor

y = x·Wᵀ where x is [rows,in] and W is stored [out,in] (HF linear convention) — computed directly, without materializing Wᵀ. Essential for big tied LM heads (avoids a huge transpose).

pub fn matmul_bt_act(&self, w: &Tensor, act: u32) -> Tensor

y = act(x·Wᵀ) — a linear projection with the activation fused into the matmul epilogue (one kernel, no intermediate). act: 0 identity, 1 relu, 2 silu, 3 gelu, 4 sigmoid. Every gated FFN (silu(x·Wgateᵀ)) and every relu/gelu MLP hidden layer collapses to a single dispatch.

pub fn gather_rows(&self, idx: &[u32]) -> Tensor

Row gather (embedding lookup): self is a [vocab, d] table; returns [idx.len(), d].

pub fn sum(&self, axes: &[usize], keepdim: bool) -> Tensor

pub fn max(&self, axes: &[usize], keepdim: bool) -> Tensor

pub fn mean(&self, axes: &[usize], keepdim: bool) -> Tensor

pub fn matmul(&self, other: &Tensor) -> Tensor

pub async fn autotune_matmul(&self, other: &Tensor) -> &'static str

Autotune the GEMM kernel for this shape on this device: time naive vs tiled and cache the winner (keyed by shape bucket). Subsequent matmuls of the same bucket use it. Returns the choice. This is how GEMM stays fast portably — the winner differs across GPUs.

pub fn matmul_tiled(&self, other: &Tensor) -> Tensor

Force the register-blocked tiled 2D GEMM (for benchmarking / large matmuls).

pub fn matmul_naive(&self, other: &Tensor) -> Tensor

The naive (non-tiled) matmul, kept for benchmarking the tiled fast-path against.

Trait Implementations

impl Clone for Tensor

fn clone(&self) -> Tensor

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.