Struct Tensor 

pub struct Tensor { /* private fields */ }

A tensor handle.

In the lazy graph model a Tensor is a lightweight reference to a node in the computation graph. Operations build up the graph; actual computation happens when eval() is called (or implicitly via to_vec_f32()).

Implementations

impl Tensor

pub fn zeros(shape: &Shape, dtype: DType, device: &Device) -> Result<Self>

Create a tensor filled with zeros.

pub fn ones(shape: &Shape, dtype: DType, device: &Device) -> Result<Self>

Create a tensor filled with ones.

pub fn from_f32(data: &[f32], shape: &Shape, device: &Device) -> Result<Self>

Create a tensor from f32 data.

pub fn from_f32_on_stream( data: &[f32], shape: &Shape, stream: &Arc<Stream>, ) -> Result<Self>

Create a tensor from f32 data on a specific stream.

pub fn from_data_with_dtype( data: Vec<f32>, shape: &Shape, dtype: DType, device: &Device, ) -> Result<Self>

Create a tensor from f32 data with a specified dtype.

The data is stored as f32 internally; dtype records the logical type (e.g. F16 weights that were converted to f32 on load).

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

Element-wise addition.

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

Element-wise subtraction.

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

Element-wise multiplication.

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

Element-wise division.

pub fn neg(&self) -> Tensor

Element-wise negation.

pub fn sum_axis(&self, axis: i32) -> Result<Tensor>

Sum along an axis.

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

Sum all elements to a scalar.

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

Matrix multiplication (2D only for now).

pub fn reshape(&self, new_shape: &Shape) -> Result<Tensor>

Reshape the tensor.

pub fn transpose(&self, axes: Option<&[usize]>) -> Result<Tensor>

Transpose (reverses axes by default, or use specified permutation).

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

Softmax along an axis.

pub fn silu(&self) -> Tensor

SiLU (Sigmoid Linear Unit) activation.

pub fn gelu(&self) -> Tensor

GELU (Gaussian Error Linear Unit) activation.

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

Layer normalization over the last dimension.

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

RMS normalization over the last dimension.

pub fn rope(&self, rotary_dim: usize, pos_offset: usize, theta: f32) -> Tensor

Apply Rotary Positional Embeddings.

pub fn layer_norm_vjp(&self, input: &Tensor, eps: f32) -> Result<Tensor>

LayerNorm VJP: compute grad_input given grad_output and original input.

pub fn rms_norm_vjp(&self, input: &Tensor, eps: f32) -> Result<Tensor>

RmsNorm VJP: compute grad_input given grad_output and original input.

pub fn softmax_vjp(&self, softmax_output: &Tensor, axis: i32) -> Result<Tensor>

Softmax VJP: compute grad_input given grad_output (self) and softmax output.

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

SiLU VJP: compute grad_input given grad_output (self) and original input.

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

GELU VJP: compute grad_input given grad_output (self) and original input.

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

Embedding lookup: gather rows from this weight matrix [vocab, dim] using indices [seq_len]. Returns [seq_len, dim].

pub fn narrow(&self, axis: i32, start: i64, length: i64) -> Result<Tensor>

Narrow (slice) along an axis: extract length elements starting at start.

pub fn cat(tensors: &[&Tensor], axis: i32) -> Result<Tensor>

Concatenate tensors along an axis.

pub fn attention( &self, k: &Tensor, v: &Tensor, scale: f32, causal: bool, ) -> Result<Tensor>

Single-head attention: Q @ K^T * scale → causal mask → softmax → @ V. Q: [Tq, Dh], K: [Tk, Dh], V: [Tk, Dh] → Output: [Tq, Dh]

pub fn sqrt(&self) -> Tensor

Element-wise square root.

pub fn eval(&self) -> Result<()>

Materialize the tensor — triggers evaluation of the computation graph.

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

Copy data out as Vec. Triggers evaluation if needed.

pub fn shape(&self) -> &Shape

Get the tensor shape.

pub fn dtype(&self) -> DType

Get the tensor dtype.

pub fn device(&self) -> &Device

Get the tensor device.

pub fn numel(&self) -> i64

Number of elements.

pub fn node_id(&self) -> NodeId

Get the graph node ID.

pub fn stream(&self) -> Arc<Stream>

Get the stream this tensor belongs to.

pub fn from_node_id( node_id: NodeId, shape: Shape, dtype: DType, device: Device, stream: Arc<Stream>, ) -> Self

Reconstruct a tensor handle from a node ID and metadata.

Used by autograd to create handles for graph introspection.

pub fn broadcast_to(&self, target: &Shape) -> Result<Tensor>

Broadcast this tensor to the target shape (numpy-style rules).

Trait Implementations

impl Add for &Tensor

type Output = Result<Tensor, MlxError>

The resulting type after applying the + operator.

fn add(self, rhs: &Tensor) -> Self::Output

Performs the + operation.

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 Mul for &Tensor

type Output = Result<Tensor, MlxError>

The resulting type after applying the * operator.

fn mul(self, rhs: &Tensor) -> Self::Output

Performs the * operation.

impl Neg for &Tensor

type Output = Tensor

The resulting type after applying the - operator.

fn neg(self) -> Self::Output

Performs the unary - operation.

impl Sub for &Tensor

type Output = Result<Tensor, MlxError>

The resulting type after applying the - operator.

fn sub(self, rhs: &Tensor) -> Self::Output

Performs the - operation.