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:
.. code-block:: pythonfrom max.experimental import tensor from max.driver import CPU from max.dtype import DType
Create and operate on tensors
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.
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
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)
argmax()
argmax()
-
Return type:
cast()
cast(dtype)
constant()
classmethod constant(value, *, dtype=None, device=None)
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:
from_graph_value()
classmethod from_graph_value(value)
-
Parameters:
-
value (TensorValue | BufferValue)
-
Return type:
full()
classmethod full(shape, value, *, dtype=None, device=None)
full_like()
classmethod full_like(type, value)
-
Parameters:
-
- type (TensorType)
- value (float | number[Any])
-
Return type:
item()
item()
max()
max()
-
Return type:
num_elements()
num_elements()
-
Return type:
ones()
classmethod ones(shape, *, dtype=None, device=None)
ones_like()
classmethod ones_like(type)
-
Parameters:
-
type (TensorType)
-
Return type:
permute()
permute(dims)
rank
property rank: int
real
property real: bool
realize
property realize
Force the tensor to realize if it is not already.
reshape()
reshape(shape)
shape
property shape: Shape
storage
Underlying memory for a realized tensor.
to()
to(device)
transpose()
transpose(dim1, dim2)
type
property type: TensorType
zeros()
classmethod zeros(shape, *, dtype=None, device=None)
zeros_like()
classmethod zeros_like(type)
-
Parameters:
-
type (TensorType)
-
Return type:
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)
driver_tensor_type()
max.experimental.tensor.driver_tensor_type(t)
-
Parameters:
-
t (Tensor)
-
Return type:
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!