Python module

tensor

Provides experimental tensor operations with eager execution capabilities.

Warning: This module contains experimental APIs that are subject to change or removal in future versions. Use with caution in production environments.

This module provides the tensor class which supports eager execution of tensor operations, complementing the graph-based execution model provided by graph. The tensor operations automatically compile and execute using the MAX runtime.

Key Features:

Eager execution: Operations execute immediately rather than building a graph.
Automatic compilation: Tensors are compiled and optimized automatically.
Lazy evaluation: Tensors may be computed lazily until their values are needed.
NumPy compatibility: Supports common NumPy-like operations and indexing.

Create and manipulate tensors with automatic compilation and optimization:

from max.experimental import tensor
from max.driver import CPU
from max.dtype import DType

x = tensor.Tensor.ones((2, 3), dtype=DType.float32, device=CPU())
y = tensor.Tensor.zeros((2, 3), dtype=DType.float32, device=CPU())
result = x + y  # Eager execution with automatic compilation

`ComputeGraph`

class max.experimental.tensor.ComputeGraph(context=None, sources=(), seed=0)

Compute graph storage for unrealized tensors.

The compute graph is a directed acyclic graph.

There is a single global compute graph we use for Tensor operations. New tensors are added as nodes to this graph by tensor operations. Once they are realized the graph is simplified and the newly realized tensors become sources of the graph.

Terminology:

A “source” of the graph is a realized tensor that some unrealized tensor depends on.
“unrealized” refers to a node in the graph which is not a source, or to the tensor object that it backs. There is a 1:1 relationship between the node and the tensor object.

It is not obvious a priori which unrealized nodes to evaluate at what time. The evaluate method of the graph is at its heart a heuristic choosing among various tradeoffs of what to compute.

The current implementation first prunes the graph of all dead nodes (nodes which no longer have live python references to them) and then realizes all remaining nodes. This is an implementation detail and is subject to change.

Parameters:

context (mlir.Context | None)
sources (dict[Value[Any], Tensor])
seed (int)

`add_source()`

add_source(tensor)

Parameters:: tensor (Tensor)
Return type:: None

`add_unrealized()`

add_unrealized(tensor)

Parameters:: tensor (Tensor)
Return type:: None

`evaluate()`

async evaluate(tensor)

Realize the input tensor object.

It is currently undefined to operate on tensors during evaluation.

After execution:

The compute graph object and all tensors realized or otherwise will be in valid states.
The input tensor is guaranteed to be realized.
Some other previously unrealized tensors may be realized
Any realized tensors with live references will not be unrealized.

Parameters:: tensor (Tensor)
Return type:: None

`graph`

graph: Graph

`sources`

sources: dict[Value[Any], Tensor]

Keeps a strong reference to tensor data that we need to compute graph values

`unrealized`

unrealized: WeakValueDictionary[int, Tensor]

Keeps weak references to intermediate unrealized tensor values, which may never need to be realized.

`Tensor`

class max.experimental.tensor.Tensor(*, storage=None, value=None)

A Tensor object with numerics.

A Tensor type that can do the kinds of things people expect tensors to do.

Tensor operations should always meet the following criteria:

Any illegal operation on a tensor must fail immediately with a python exception with a clear error message
All operations on tensors that read or write Tensor memory values use our high-performance compiler and Mojo kernel library.

The out of the box experience should be the best one available for working with Tensors and numerics, and give seemless access to direct low-level programmability in Mojo.

Notably Tensor does not require that it is backed by memory. If no side-effecting operation has been done on a Tensor object, then there is no guarantee it has been computed yet. Critically a user should never know or care whether the tensor is backed by data: the behavior should be exactly as if it were.

For discussion purposes, a “realized” tensor is a tensor which references concrete memory, and an “unrealized” one does not. An “unrealized” tensor may still have a driver tensor as storage, but this memory may not be an up-to-date reference of the tensor’s data, for instance in the case of mutating ops.

Tensors unify the graph concepts of TensorValue and BufferValue. Given x: Tensor:

If x is realized, it will be backed by a BufferValue input to the graph.
If x is unrealized:
- It will be backed by a BufferValue if the op that created it returned a buffer _or_ if BufferValue(x) is ever called
- Otherwise it will be backed by a TensorValue
x may _always_ be loaded into a TensorValue via TensorValue(x):
- If x is backed by a BufferValue, this will do a “load”. The load is for operation ordering and will be optimized away.
x may _always_ be loaded into a BufferValue via BufferValue(x):
- Afterwards, x will always be unrealized, since it may now have been passed into side-effecting ops.
- If x was backed by a TensorValue, it will now be backed by a new BufferValue from ops.buffer_create containing the same data.

This allows Tensors to be transparently treated as being mutable buffers or immutable tensors while encoding only the necessary semantics into the compilation graph, and mutating computations remain lazy.

Parameters:

storage (Tensor | None)
value (graph.BufferValue | graph.TensorValue | None)

`T`

property T: Tensor

`arange()`

classmethod arange(start=0, stop=None, step=1, *, dtype=None, device=None)

Parameters:

start (int)
stop (int | None)
step (int)
dtype (DType | None)
device (Device | None)

Return type:

Tensor

`argmax()`

argmax(axis=-1)

Parameters:: axis (int)
Return type:: Tensor

`cast()`

cast(dtype)

Parameters:: dtype (DType)
Return type:: Tensor

`constant()`

classmethod constant(value, *, dtype=None, device=None)

Creates a constant tensor from a scalar, array, or nested list.

Constructs a tensor with constant values that can be a scalar, a nested Python list, or a DLPack-compatible array. The shape is automatically inferred from the input data structure.

from max.experimental import tensor
from max.dtype import DType

# Create from scalar
x = tensor.Tensor.constant(42, dtype=DType.int32)

# Create from nested list
y = tensor.Tensor.constant([[1.0, 2.0], [3.0, 4.0]])

# Create from NumPy array
import numpy as np

z = tensor.Tensor.constant(np.array([1, 2, 3]))

Parameters:

value (DLPackArray | Sequence[float | number[Any] | Sequence[Number | NestedArray]] | float | number[Any]) – The constant value for the tensor. Can be a scalar number, a nested Python list, or any DLPack-compatible array.
dtype (DType | None) – The data type for the tensor elements. If not specified, defaults to DType.float32 for CPU devices and DType.bfloat16 for accelerator devices.
device (Device | None) – The device where the tensor will be allocated. If not specified, defaults to an accelerator if available, otherwise CPU.

Returns:

A new tensor containing the constant value(s).

Return type:

Tensor

`device`

property device: Device

The tensor’s device.

`driver_tensor`

property driver_tensor: Tensor

A pointer to the underlying memory.

Raises if the tensor is unrealized.

`dtype`

property dtype: DType

`from_dlpack()`

classmethod from_dlpack(array)

Creates a tensor from a DLPack array.

Constructs a tensor by importing data from any object that supports the DLPack protocol (such as NumPy arrays and PyTorch tensors). This enables zero-copy interoperability with other array libraries.

import numpy as np
from max.experimental import tensor

# Create a NumPy array
np_array = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)

# Convert to MAX tensor via DLPack
x = tensor.Tensor.from_dlpack(np_array)

Parameters:: array (DLPackArray) – Any object supporting the DLPack protocol, such as NumPy arrays, PyTorch tensors, or JAX arrays.
Returns:: A new tensor containing the data from the DLPack array.
Return type:: Tensor

`from_graph_value()`

classmethod from_graph_value(value)

Creates a tensor from a graph value.

Constructs a tensor from an existing graph value, which can be either a TensorValue or BufferValue. This is useful for converting graph level values into tensor objects for eager execution.

Parameters:: value (TensorValue | BufferValue) – The graph value to wrap. Can be either a TensorValue or BufferValue from the MAX graph API.
Returns:: A new tensor backed by the provided graph value.
Return type:: Tensor

`full()`

classmethod full(shape, value, *, dtype=None, device=None)

Creates a tensor filled with a specified value.

Returns a new tensor with the given shape where all elements are initialized to the specified value. This is useful for creating tensors with uniform values other than zero or one.

from max.experimental import tensor
from max.dtype import DType

# Create a 3x3 tensor filled with 7
x = tensor.Tensor.full((3, 3), value=7, dtype=DType.int32)

# Create a 2x4 tensor filled with pi
y = tensor.Tensor.full((2, 4), value=3.14159)

Parameters:

shape (Iterable[int | str | Dim | integer[Any]]) – The shape of the output tensor. Can be a tuple of integers, a list of integers, or any value that can be converted to a shape.
value (float | number[Any]) – The scalar value to fill the tensor with.
dtype (DType | None) – The data type for the tensor elements. If not specified, defaults to DType.float32 for CPU devices and DType.bfloat16 for accelerator devices.
device (Device | None) – The device where the tensor will be allocated. If not specified, defaults to an accelerator if available, otherwise CPU.

Returns:

A new tensor with the specified shape filled with the given value.

Return type:

Tensor

`full_like()`

classmethod full_like(type, value)

Creates a tensor filled with a value, matching a given type’s properties.

Returns a new tensor filled with the specified value that matches the shape, data type, and device of the given tensor type. This is useful when you need to create a tensor with uniform values that’s compatible with an existing tensor’s properties.

from max.experimental import tensor
from max.graph import TensorType
from max.driver import CPU
from max.dtype import DType

# Create a reference tensor type
ref_type = TensorType(DType.float32, (2, 3), device=CPU())

# Create tensor filled with 5.0 matching the reference type
x = tensor.Tensor.full_like(ref_type, value=5.0)

Parameters:

type (TensorType) – The tensor type to match. The returned tensor will have the same shape, dtype, and device as this type.
value (float | number[Any]) – The scalar value to fill the tensor with.

Returns:

A new tensor filled with the specified value, matching the properties of the input type.

Return type:

Tensor

`item()`

item()

`max()`

max(axis=-1)

Parameters:: axis (int)
Return type:: Tensor

`mean()`

mean(axis=-1)

Parameters:: axis (int)
Return type:: Tensor

`num_elements()`

num_elements()

Return type:: int

`ones()`

classmethod ones(shape, *, dtype=None, device=None)

Parameters:

shape (Iterable[int | str | Dim | integer[Any]])
dtype (DType | None)
device (Device | None)

Return type:

Tensor

`ones_like()`

classmethod ones_like(type)

Parameters:: type (TensorType)
Return type:: Tensor

`permute()`

permute(dims)

Parameters:: dims (list[int])
Return type:: Tensor

`range_like()`

classmethod range_like(type)

Parameters:: type (TensorType)
Return type:: Tensor

`rank`

property rank: int

`real`

property real: bool

`realize`

property realize

Force the tensor to realize if it is not already.

`reshape()`

reshape(shape)

Parameters:: shape (Iterable[int | str | Dim | integer[Any]])
Return type:: Tensor

`shape`

property shape: Shape

`storage`

storage: Tensor | None

Underlying memory for a realized tensor.

`to()`

to(device)

Parameters:: device (Device)
Return type:: Tensor

`transpose()`

transpose(dim1, dim2)

Parameters:

dim1 (int)
dim2 (int)

Return type:

Tensor

`type`

property type: TensorType

`zeros()`

classmethod zeros(shape, *, dtype=None, device=None)

Creates a tensor filled with zeros.

Returns a new tensor with the specified shape where all elements are initialized to zero. The tensor is created with eager execution and automatic compilation.

from max.experimental import tensor
from max.driver import CPU
from max.dtype import DType

# Create a 2x3 tensor of zeros
x = tensor.Tensor.zeros((2, 3), dtype=DType.float32, device=CPU())
# Result: [[0.0, 0.0, 0.0],
#          [0.0, 0.0, 0.0]]

# Create a 1D tensor using default dtype and device
y = tensor.Tensor.zeros((5,))

Parameters:

shape (Iterable[int | str | Dim | integer[Any]]) – The shape of the output tensor. Can be a tuple of integers, a list of integers, or any value that can be converted to a shape.
dtype (DType | None) – The data type for the tensor elements. If not specified, defaults to DType.float32 for CPU devices and DType.bfloat16 for accelerator devices.
device (Device | None) – The device where the tensor will be allocated. If not specified, defaults to an accelerator if available, otherwise CPU.

Returns:

A new tensor with the specified shape filled with zeros.

Return type:

Tensor

`zeros_like()`

classmethod zeros_like(type)

Creates a tensor of zeros matching a given type’s properties.

Returns a new tensor filled with zeros that matches the shape, data type, and device of the specified tensor type. This is useful when you need to create a zero tensor that’s compatible with an existing tensor’s properties.

from max.experimental import tensor
from max.graph import TensorType
from max.driver import CPU
from max.dtype import DType

# Create a reference tensor type
ref_type = TensorType(DType.float32, (3, 4), device=CPU())

# Create zeros tensor matching the reference type
x = tensor.Tensor.zeros_like(ref_type)
# Result: 3x4 tensor of zeros with dtype float32 on CPU

Parameters:: type (TensorType) – The tensor type to match. The returned tensor will have the same shape, dtype, and device as this type.
Returns:: A new tensor filled with zeros matching the properties of the input type.
Return type:: Tensor

`contextvar_context()`

max.experimental.tensor.contextvar_context(var, value)

Parameters:

var (ContextVar[T])
value (T)

`default_device()`

max.experimental.tensor.default_device(device)

Context manager for setting the default device for tensors.

Parameters:: device (Device)

`default_dtype()`

max.experimental.tensor.default_dtype(dtype)

Context manager for setting the default dtype for tensors.

Parameters:: dtype (DType)

`defaults()`

max.experimental.tensor.defaults(dtype=None, device=None)

Parameters:

dtype (DType | None)
device (Device | None)

Return type:

tuple[DType, Device]

`driver_tensor_type()`

max.experimental.tensor.driver_tensor_type(t)

Parameters:: t (Tensor)
Return type:: TensorType

ComputeGraph​

add_source()​

add_unrealized()​

evaluate()​

graph​

sources​

unrealized​

Tensor​

T​

arange()​

argmax()​

cast()​

constant()​

device​

driver_tensor​

dtype​

from_dlpack()​

from_graph_value()​

full()​

full_like()​

item()​

max()​

mean()​

num_elements()​

ones()​

ones_like()​

permute()​

range_like()​

rank​

real​

realize​

reshape()​

shape​

storage​

to()​

transpose()​

type​

zeros()​

zeros_like()​

contextvar_context()​

default_device()​

default_dtype()​

defaults()​

driver_tensor_type()​

`ComputeGraph`

`add_source()`

`add_unrealized()`

`evaluate()`

`graph`

`sources`

`unrealized`

`Tensor`

`T`

`arange()`

`argmax()`

`cast()`

`constant()`

`device`

`driver_tensor`

`dtype`

`from_dlpack()`

`from_graph_value()`

`full()`

`full_like()`

`item()`

`max()`

`mean()`

`num_elements()`

`ones()`

`ones_like()`

`permute()`

`range_like()`

`rank`

`real`

`realize`

`reshape()`

`shape`

`storage`

`to()`

`transpose()`

`type`

`zeros()`

`zeros_like()`

`contextvar_context()`

`default_device()`

`default_dtype()`

`defaults()`

`driver_tensor_type()`