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

`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: WeakSet[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.

Parameters:

storage (driver.Tensor | None)
value (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()

Return type:: Tensor

`cast()`

cast(dtype)

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

`constant()`

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

Parameters:

value (ndarray[Any, dtype[number[Any]]] | int | float | Sequence[int | float])
dtype (DType | None)
device (Device | None)

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)

Parameters:: array (DLPackArray)
Return type:: Tensor

`from_tensor_value()`

classmethod from_tensor_value(value)

Parameters:: value (TensorValue)
Return type:: Tensor

`full()`

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

Parameters:

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

Return type:

Tensor

`full_like()`

classmethod full_like(type, value)

Parameters:

type (TensorType)
value (int | float)

Return type:

Tensor

`item()`

item()

`max()`

max()

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

`rank`

property rank: int

`real`

property real: bool

Whether the tensor is realized or not.

`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

`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)

Parameters:

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

Return type:

Tensor

`zeros_like()`

classmethod zeros_like(type)

Parameters:: type (TensorType)
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

WARNING​

ComputeGraph​

add_source()​

evaluate()​

graph​

sources​

unrealized​

Tensor​

T​

arange()​

argmax()​

cast()​

constant()​

device​

driver_tensor​

dtype​

from_dlpack()​

from_tensor_value()​

full()​

full_like()​

item()​

max()​

num_elements()​

ones()​

ones_like()​

permute()​

rank​

real​

realize​

reshape()​

shape​

to()​

transpose()​

type​

zeros()​

zeros_like()​

contextvar_context()​

default_device()​

default_dtype()​

defaults()​

driver_tensor_type()​

WARNING

`ComputeGraph`

`add_source()`

`evaluate()`

`graph`

`sources`

`unrealized`

`Tensor`

`T`

`arange()`

`argmax()`

`cast()`

`constant()`

`device`

`driver_tensor`

`dtype`

`from_dlpack()`

`from_tensor_value()`

`full()`

`full_like()`

`item()`

`max()`

`num_elements()`

`ones()`

`ones_like()`

`permute()`

`rank`

`real`

`realize`

`reshape()`

`shape`

`to()`

`transpose()`

`type`

`zeros()`

`zeros_like()`

`contextvar_context()`

`default_device()`

`default_dtype()`

`defaults()`

`driver_tensor_type()`