Python module

max.experimental.tensor

Provides tensor operations with eager execution capabilities.

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 semantics: Operations give immediate results for quick iteration and feedback.
High performance: All operations use high-performance Mojo implementations compiled specifically for the available hardware.
Automatic compilation: Tensors are compiled and optimized automatically. Operations may be easily fused into larger graphs to take advantage of the graph compiler’s automatic fusions.
Lazy evaluation: Tensors may be computed lazily until their values are needed.
Familiar API: Supports common array operations and indexing.

Note

Tensors use lazy evaluation and JIT compilation, which incurs compilation overhead on first execution. This can result in higher latency for initial operations compared to eager frameworks like NumPy or PyTorch. Subsequent executions reuse compiled kernels for better performance.

Create and manipulate tensors with automatic compilation and optimization:

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

# Create and operate on tensors
x = Tensor.ones((2, 3), dtype=DType.float32, device=CPU())
y = Tensor.zeros_like(x)
result = x + y  # Eager execution with automatic compilation

Operations may be combined into a single execution graph to take advantage of automatic kernel fusion:

from max.experimental import functional as F

@F.functional
def linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
    return x @ weight.T + bias

# Create and operate on tensors
x = Tensor.ones([2, 3])
weight = Tensor.ones([6, 3])
bias = Tensor.ones([6])

# Eager execution with a single fused graph
result = linear(x, weight, bias)

Users may opt in to lazy execution. This is primarily useful for

Operations which may never execute, for instance creating modules with randomly initialized weights before loading weights
Combining many operations into a single execution

from max.experimental.nn import Linear

with F.lazy():
    model = Linear(2, 3)

print(model)  # Lazy weights not initialized

# Load pretrained weights
weights =  {
    "weight": Tensor.zeros([3, 2]),
    "bias": Tensor.zeros([3]),
}
model.load_state_dict(weights)

# Or compile directly without ever initializing weights
from max.graph import TensorType
input_type = TensorType(DType.float32, ["batch", 2], CPU())
model = model.compile(input_type, weights=weights)

Tensor

`Tensor`	A multi-dimensional array with eager execution and automatic compilation.

Defaults

`default_device`	Context manager for setting the default device for tensor creation.
`default_dtype`	Context manager for setting the default dtype for tensor creation.
`defaults`	Gets the default dtype and device for tensor creation.
`defaults_like`	Context manager setting the default dtype and device for tensor creation.

Realization

`RealizationContext`	Implements a way to realize unrealized tensors.
`RealizationState`	State for an unrealized tensor.

`current_realization_context`	Return a value for the context variable for the current context.
`realization_context`	Sets the current realization context, within a context manager.

Tensor​

Defaults​

Realization​

Tensor

Defaults

Realization