Struct Tensor 

pub struct Tensor { /* private fields */ }

A CPU tensor: a Shape plus a strided view into a shared, ref-counted Storage buffer.

Why Arc<Storage> plus separate strides and an offset

Model weights are hundreds of megabytes; every op that produces a new logical view of existing data — reshape, transpose, narrow, broadcast_to — must not copy that data. Sharing storage via Arc makes Tensor::clone() and every pure view operation an O(1) pointer bump plus a small Shape/Vec<usize> copy, never an O(n) memcpy of element data. strides (in elements, one per dimension) and offset (in elements, into the shared buffer) are what let a view disagree with its storage about layout: narrow changes offset, transpose reorders strides, broadcast_to introduces zero strides.

Why quantized tensors don’t get real views

A DType::is_quantized tensor’s strides are always its shape’s canonical strides and its offset is always 0 — block-quantized data cannot be addressed at anything finer than a whole block, so “a strided view into the middle of a Q4_0 tensor” is not a representable concept. Every op that would need one (Tensor::narrow, Tensor::transpose, Tensor::broadcast_to, the arithmetic and normalization ops) rejects quantized tensors with Error::QuantizedElementAccess via Tensor::require_elementwise and Tensor::require_dtype. Call Tensor::to_dtype with DType::F32 first.

Implementations

impl Tensor

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

Batched matrix multiplication: self is [..batch, m, k], other is [..batch, k, n], and the result is [..batch, m, n].

Both operands must be rank >= 2 and f32. The leading (“batch”) dimensions are broadcast against each other with the same right-aligned rule as kopitiam_core::Shape::broadcast (so a single [k, n] weight matrix multiplies against a [batch, m, k] activation without the caller manually replicating it) — this is what makes a single implementation handle both the plain 2D case (empty batch shape) and the batched case identically, rather than special-casing 2D.

Errors

Shape problems (rank < 2, or the inner dimension k disagreeing between operands) are reported as Error::ShapeMismatch. That variant’s field names (expected/actual) were designed for reshape; here they are reused to mean “these two shapes are not compatible for matrix multiplication” — the closest fit among kopitiam-core’s error variants, which this crate cannot modify. Incompatible batch dimensions surface as Error::NotBroadcastable, propagated directly from Shape::broadcast.

pub fn quantized_matmul(&self, weight: &Tensor) -> Result<Tensor>

self @ weight^T — the same semantics [crate::linear::linear]’s (in kopitiam-runtime) transpose-then-Tensor::matmul gives an f32 weight — but for a block-quantized weight (DType::Q4_0 or DType::Q8_0), computed without ever dequantizing weight to f32.

The algorithm: quantize the activation, dot in integer space

The “correct before fast” version of a quantized matmul is weight.to_dtype(DType::F32) once, then plain Tensor::matmul — that remains this crate’s default and its permanent reference implementation (see “Correctness” below). It is also why a 7B-parameter Q4_0 model, genuinely about 4GB on disk, balloons to roughly 28GB of resident f32 the instant it is loaded that way: dequantizing trades away the entire point of shipping a quantized model (fitting in memory) for arithmetic convenience.

This method is the technique ggml/llama.cpp use instead — the reason quantized inference is fast, not merely small — and it never produces an f32 copy of weight at all. Every row of self is quantized to Q8_0 blocks on the fly ([quant::quantize_row_q8_0]), and each output element becomes a per-block integer dot product ([quant::q4_0_dot_q8_0] / [quant::q8_0_dot_q8_0]) between that freshly-quantized activation block and weight‘s already-quantized block, with both sides’ scales multiplied in once at the very end of each block. weight’s bytes are read exactly as they sit in Storage::Quantized — nothing about weight is ever expanded to a wider type. This is a genuinely different algorithm from “dequantize then multiply”, not an optimization of it: the former spends 2 * out_features * in_features f32 multiply-adds; this spends that many i32 multiply-adds (integer arithmetic a CPU — including a phone-class one — does at several times f32’s throughput) plus only 2 * out_features * (in_features / 32) f32 multiplies for the per-block scale combination.

self is [..., in_features], f32, any rank >= 1; weight is [out_features, in_features], exactly rank 2, quantized. Returns [..., out_features]. in_features must be a whole number of 32-element blocks (true of every real GGUF Q4_0/Q8_0 export, since a block may never straddle a row boundary).

Correctness

This method’s own tests assert, for both Q4_0 and Q8_0, that its output agrees with weight.to_dtype(DType::F32) followed by Tensor::matmul within the error the activation’s Q8_0 quantization alone can introduce (both paths read the identical weight bits, so any difference comes only from self going through [quant::quantize_row_q8_0] on this path and not on the reference path) — see this module’s quantized_matmul_agrees_with_... tests. That reference path is why dequantizing to f32 stays in this crate permanently: it is not just a fallback, it is the oracle this method is checked against.

Errors

Error::DTypeMismatch if self is not f32. Error::UnsupportedDType if weight’s dtype is not one of the two formats this method has a fused kernel for. The other quantized formats — Q4_1/Q5_0/Q5_1 and every K-quant (Q4_K/Q5_K/Q6_K/…) — dequantize fine via Tensor::to_dtype but have no fused path here yet, so a caller wanting to matmul against one must dequantize it to f32 first; see the parent crate’s Phase 2 scope. Also errors if weight is not quantized at all (use Tensor::matmul directly for a plain f32 weight). Error::ShapeMismatch if weight is not rank 2, if self’s last dimension does not equal weight’s in_features, or if in_features is not a multiple of 32.

impl Tensor

pub fn add(&self, other: &Tensor) -> Result<Tensor>

Elementwise self + other, broadcasting shapes per kopitiam_core::Shape::broadcast.

pub fn sub(&self, other: &Tensor) -> Result<Tensor>

pub fn mul(&self, other: &Tensor) -> Result<Tensor>

pub fn div(&self, other: &Tensor) -> Result<Tensor>

Elementwise self / other. Division by zero follows ordinary IEEE 754 float semantics (+/-inf or NaN) rather than erroring — a tensor op has no way to know whether that is a bug or an expected edge case (e.g. a masked-out attention score), so it is left to the caller to detect if it matters.

impl Tensor

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

Softmax along axis: for every slice along that axis, out[i] = exp(x[i] - max(x)) / sum(exp(x - max(x))).

Why subtract the max before exp

Real transformer logits routinely reach the hundreds or thousands in magnitude. exp(1000.0) overflows f32 to +inf, and inf / inf is NaN — softmax computed the naive way silently poisons the rest of the forward pass on real inputs, not just adversarial ones. Subtracting the row’s max first is mathematically a no-op (it cancels between numerator and denominator) but keeps every exponent <= 0, so the largest possible input to exp is exp(0) = 1. This is why the max-subtraction step is not an optional optimization here — see the large_magnitude_input_does_not_overflow test below for the exact failure this avoids.

impl Tensor

pub fn sum(&self, axis: usize, keepdim: bool) -> Result<Tensor>

Sums along axis. If keepdim, axis becomes length 1 in the result instead of being removed (useful for immediately broadcasting the result back against the input, e.g. normalizing by a row sum).

pub fn max(&self, axis: usize, keepdim: bool) -> Result<Tensor>

Reduces along axis by taking the maximum.

pub fn argmax(&self, axis: usize, keepdim: bool) -> Result<Tensor>

Index of the maximum along axis, returned as an DType::I32 tensor of positions (not the max values — that’s Tensor::max).

This is the last step of greedy decoding: run the forward pass, get the final [.., vocab] logits, then argmax(last_axis, false) gives you the chosen token id per position. Read the ids out with Tensor::to_vec_i32.

The output dtype is I32 on purpose — these are indices into a dimension, the same currency Tensor::gather_rows consumes as token ids, so the next embedding lookup can eat this straight without a cast. The result’s shape is self’s shape with axis dropped, or set to length 1 if keepdim (same rule as Tensor::sum).

Tie-break contract (load-bearing, so it’s deterministic): on a tie the first (lowest-index) maximum wins. We only overwrite the running best on a strict >, never on ==, so re-running on the same input always yields the same ids. NaN never wins a comparison (v > best is false for NaN), so a NaN element is skipped rather than selected — don’t lean on that for correctness though; NaN logits mean the forward pass already went wrong upstream.

impl Tensor

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

RMSNorm (Zhang & Sennrich, 2019): x / rms(x) * weight, where rms(x) = sqrt(mean(x^2) + eps), computed independently for every vector along the last dimension.

weight must have exactly hidden elements, where hidden is self’s last dimension.

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

Layer normalization (Ba, Kiros & Hinton, 2016): (x - mean(x)) / sqrt(var(x) + eps) * weight + bias, with mean and the (population, i.e. divide-by-N) var computed independently for every vector along the last dimension.

impl Tensor

pub fn silu(&self) -> Result<Tensor>

SiLU / swish (Elfwing, Uchibe & Doya, 2017; also Hendrycks & Gimpel): x * sigmoid(x) = x / (1 + exp(-x)). Used by LLaMA-family FFNs.

pub fn gelu(&self) -> Result<Tensor>

GELU (Hendrycks & Gimpel, 2016), tanh approximation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))).

This crate implements the tanh approximation rather than the exact 0.5 * x * (1 + erf(x / sqrt(2))) form because erf is not in Rust’s std — computing it exactly would mean shipping a rational or Chebyshev erf approximation of our own, which is more code and more numerical risk than the tanh form for a difference that is below 1e-3 everywhere and is what GPT-2/GPT-NeoX-family models were themselves trained and evaluated with (the tanh form is sometimes called “gelu_new” in that lineage). If a model that specifically requires the exact erf form shows up, add it as a second method rather than replacing this one.

pub fn tanh(&self) -> Result<Tensor>

Hyperbolic tangent, elementwise: tanh(x), in (-1, 1). The cell-candidate (g) nonlinearity and the output squashing (tanh(c')) of an LSTM cell — see Tensor::lstm_cell.

Matches Tesseract’s Tanh (functions.cpp), computed directly via f32::tanh rather than through that file’s 8k-entry lookup table (TanhTable): standard math, so there is nothing Tesseract-specific to port here — the LUT is a speed/accuracy trade this crate has no reason to reproduce.

pub fn sigmoid(&self) -> Result<Tensor>

Logistic sigmoid, elementwise: 1 / (1 + exp(-x)), in (0, 1). The input/forget/output gate nonlinearity of an LSTM cell — see Tensor::lstm_cell.

Matches Tesseract’s Logistic (functions.cpp), computed directly rather than via its 8k-entry lookup table (LogisticTable); standard math, no Tesseract-specific code to port. Tensor::logistic is an alias for callers that prefer Tesseract’s name for the same function.

pub fn logistic(&self) -> Result<Tensor>

Alias for Tensor::sigmoid under Tesseract’s name for the function.

impl Tensor

pub fn lstm_cell( input_gate: &Tensor, forget_gate: &Tensor, cell_candidate: &Tensor, output_gate: &Tensor, cell_prev: &Tensor, ) -> Result<LstmState>

Computes one LSTM timestep from the four gate pre-activations and the previous cell state, returning the new LstmState.

Each argument is the pre-activation of one gate — the summed input and recurrent linear contributions, before any nonlinearity — so this helper composes cleanly after the recognizer’s gate matmuls. The activations (sigmoid on the three gates, tanh on the candidate and on the squashed state) are applied here.

All five tensors must be f32 and share one shape (the hidden size, optionally with leading batch/time dims); the result carries that same shape. Broadcasting is intentionally not offered: the gates and the state of a single LSTM layer are always the same width, so a shape disagreement is a bug to surface, not a broadcast to paper over.

Errors

Error::DTypeMismatch if any operand is not f32 (propagated from the underlying Tensor::sigmoid/Tensor::tanh/Tensor::mul). Error::ShapeMismatch if the operands’ shapes are not all identical.

impl Tensor

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

1-D convolution (cross-correlation, the ML convention — the kernel is not flipped), no padding (“valid”), with a configurable stride.

self is the input [in_channels, length]; weight is [out_channels, in_channels, kernel]; bias, if given, is [out_channels]. The result is [out_channels, out_length] where out_length = (length - kernel) / stride + 1. Each output element is

out[oc][t] = bias[oc] + sum_ic sum_k input[ic][t*stride + k] * weight[oc][ic][k]

Errors

Error::DTypeMismatch if any operand is not f32. Error::ShapeMismatch if self is not rank 2, weight is not rank 3, the channel counts disagree, stride is 0, or kernel is larger than length (which would make out_length empty or ill-defined).

pub fn maxpool1d(&self, kernel: usize, stride: usize) -> Result<Tensor>

1-D max-pooling over a sliding window, no padding, with a configurable stride. Applied independently per channel.

self is [channels, length]; the result is [channels, out_length] where out_length = (length - kernel) / stride + 1 and out[c][t] = max(input[c][t*stride .. t*stride + kernel]).

Errors

Error::DTypeMismatch if self is not f32. Error::ShapeMismatch if self is not rank 2, stride or kernel is 0, or kernel exceeds length.

impl Tensor

pub fn tessdata_int8_to_f32( weights: &[i8], scales: &[f32], rows: usize, cols: usize, ) -> Result<Tensor>

Decodes a Tesseract int8 LSTM weight matrix to an f32 Tensor.

weights is the row-major [rows, cols] int8 matrix (rows = num_outputs, cols = num_inputs), and scales is the per-row recovery scale — one entry per row, the value for which w ~= weights[i][j] * scales[i] (Tesseract’s on-disk max_abs / 127; see the module docs for the two-scaling subtlety). Returns the dequantized [rows, cols] tensor out[i][j] = weights[i*cols + j] as f32 * scales[i], ready for Tensor::matmul.

Errors

Error::ShapeMismatch if weights.len() != rows * cols or scales.len() != rows.

impl Tensor

pub fn gather_rows(&self, indices: &Tensor) -> Result<Tensor>

Gathers rows of a [vocab, hidden] embedding table by token id.

indices (dtype DType::I32, any shape) holds token ids in [0, vocab). The result has shape indices.shape() + [hidden] — every index becomes one full row of self.

Why this is not torch.gather

PyTorch’s gather(dim, index) requires index to have the same rank as the source and picks one scalar per output position along dim. That is a general primitive this crate deliberately does not implement: nothing in a forward pass needs it. What every transformer does need is embedding lookup — pick whole rows out of a [vocab, hidden] table by a batch of token ids — which is a different (simpler, more common) shape of operation with its own name here rather than a special case of a more general gather that this crate does not otherwise use.

impl Tensor

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

Concatenates tensors along dim. All tensors must share a dtype (any non-quantized dtype — this is pure data movement, not arithmetic, so it is not restricted to f32 the way the math ops are) and must agree on every dimension except dim.

impl Tensor

pub fn from_f32(data: Vec<f32>, shape: impl Into<Shape>) -> Result<Self>

Builds a contiguous f32 tensor from data, which must hold exactly shape.elem_count() elements.

pub fn from_f16(data: Vec<u16>, shape: impl Into<Shape>) -> Result<Self>

Builds a tensor from raw IEEE 754 binary16 bits. Values are decoded lazily by ops via Tensor::to_dtype; nothing here interprets them.

pub fn from_bf16(data: Vec<u16>, shape: impl Into<Shape>) -> Result<Self>

Builds a tensor from raw bfloat16 bits.

pub fn from_i8(data: Vec<i8>, shape: impl Into<Shape>) -> Result<Self>

pub fn from_i32(data: Vec<i32>, shape: impl Into<Shape>) -> Result<Self>

Builds an i32 tensor — the dtype kopitiam_core::DType documents as “mostly for token ids and indices”, e.g. the input to Tensor::gather_rows.

pub fn from_quantized( dtype: DType, bytes: Vec<u8>, shape: impl Into<Shape>, ) -> Result<Self>

Builds a block-quantized tensor from raw on-disk block bytes.

bytes must be exactly the number of bytes shape.elem_count() elements of dtype require; see Storage::new_quantized for the precise failure modes (Error::PartialQuantizedBlock, Error::StorageTooSmall).

pub fn zeros(shape: impl Into<Shape>) -> Self

A tensor of f32 zeros. Convenience for tests and for initializing accumulators; not meant as a general “any dtype” factory, since every op in this crate computes in f32 (see the crate docs).

pub fn shape(&self) -> &Shape

pub fn dtype(&self) -> DType

pub fn device(&self) -> Device

pub fn rank(&self) -> usize

pub fn elem_count(&self) -> usize

pub fn strides(&self) -> &[usize]

This view’s element strides. Not necessarily the canonical strides of shape() — see Tensor::is_contiguous.

pub fn is_contiguous(&self) -> bool

Whether this view’s elements sit sequentially in storage in row-major order with no gaps, i.e. whether strides() equals shape().strides(). A fresh tensor from any from_* constructor is always contiguous; narrow, transpose, and broadcast_to can all produce non-contiguous views.

pub fn to_vec_f32(&self) -> Result<Vec<f32>>

Copies this view’s elements out as a plain Vec<f32> in row-major order. Requires dtype() == DType::F32; call Tensor::to_dtype first for any other dtype, including the quantized ones.

pub fn to_vec_i32(&self) -> Result<Vec<i32>>

Copies this view’s elements out as a plain Vec<i32> in row-major order. The i32 twin of Tensor::to_vec_f32: needs dtype() == DType::I32, errors Error::DTypeMismatch otherwise.

This is how you read the output of an index-producing op — namely Tensor::argmax, whose result is the token ids a greedy sampler hands back to the runtime. Without it argmax would be a dead end: you could compute the winning indices but never get them out as plain numbers. (There’s deliberately no i32 <-> f32 cast anywhere in this crate — see Tensor::to_dtype — so this accessor is the only way I32 data leaves a tensor, same as to_vec_f32 is for F32.)

pub fn contiguous(&self) -> Result<Tensor>

Returns a tensor equal to self but guaranteed contiguous, copying only if self is not already contiguous (in which case it is a cheap Arc clone).

This is the shared implementation behind Tensor::reshape and Tensor::concat: both need row-major-sequential data to work with and neither cares whether that came for free or required a copy.

pub fn to_dtype(&self, dtype: DType) -> Result<Tensor>

Converts to dtype, decoding as needed.

Converting to dtype() == dtype is a cheap Arc clone. Converting to DType::F32 is supported from every dtype this crate knows, including every quantized format (this is the dequantization entry point — see [crate::quant] for the block layouts) and f16/bf16 (see [crate::half]). Converting from f32 to f16 or bf16 is also supported, for storing activations compactly.

Not supported: converting f32 back to a quantized format. Requantization is a model-export/training-time concern — it requires choosing a quantization scheme (per-tensor vs per-block scale, calibration, error-diffusion) — not a forward-pass inference concern, which is this crate’s whole scope. i8/i32 casts to or from f32 are likewise not implemented: nothing in the ops this crate provides needs them (token ids are consumed as i32 directly by Tensor::gather_rows, never mixed arithmetically with f32).

pub fn transpose(&self, dim0: usize, dim1: usize) -> Result<Tensor>

Returns a view with dim0 and dim1 swapped. Zero-copy: only the shape and strides change.

pub fn narrow(&self, dim: usize, start: usize, len: usize) -> Result<Tensor>

Returns the sub-view [start, start + len) along dim. Zero-copy: only the shape and offset change.

pub fn broadcast_to(&self, shape: impl Into<Shape>) -> Result<Tensor>

Returns a view broadcast to shape, following Shape::broadcast’s NumPy right-aligned rule. Zero-copy: dimensions being stretched from size 1 get stride 0, so every “broadcast” element reads the same underlying value.

pub fn reshape(&self, dims: impl Into<Vec<usize>>) -> Result<Tensor>

Reinterprets this tensor as dims, which must describe the same number of elements. Zero-copy when self is already contiguous; otherwise materializes a contiguous copy first (see Tensor::contiguous).

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.

impl Debug for Tensor

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.