IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

Tensor

Tensor​

class max.experimental.tensor.Tensor(data=None, *, dtype=None, device=None, storage=None, state=None)

source

Bases: DLPackArray, HasTensorValue

A multi-dimensional array of numeric values on a CPU or accelerator device.

You can create tensors using:

Tensors support the DLPack protocol for zero-copy data exchange with NumPy, PyTorch, JAX, and other array libraries.

import numpy as np
from max.experimental.tensor import Tensor
from max.dtype import DType

x = Tensor(42, dtype=DType.int32)
y = Tensor([[1.0, 2.0], [3.0, 4.0]])

# Create from any DLPack-compatible array; dtype is inherited
z = Tensor(np.array([1, 2, 3], dtype=np.int16))

# Use factory methods like ones, zeros, arange
zeros = Tensor.zeros((2, 2))

# Compute with Python operators or the functional API
result = y + zeros
print(result)

A tensor can span several devices. It reports a global shape folded from the per-device shapes, and each device holds its own value: a slice, a full copy, one term of a pending sum, or, on an Unknown axis, a value with no known relation to its peers. Its DeviceMapping records which. Ops accept a distributed tensor the way they accept any other tensor, and reshard it when their sharding rule requires a different placement. These accessors read one local buffer and raise ValueError on a distributed tensor: storage, state, device, driver_tensor, __dlpack__, __dlpack_device__, and bool(tensor). Awaiting the tensor raises the same error, and type raises TypeError and points at layout. item() raises unless every mesh axis is replicated. Use local_shards for the per-device values, and materialize() or to_numpy() to gather the whole array into one.

A tensor is either realized (backed by a concrete Buffer in memory) or unrealized (backed by a symbolic graph value). MAX realizes tensors by running pre-compiled Mojo kernels or JIT-compiled graphs.

Parameters:

  • data (DLPackArray | NestedArray | Number | None) – The value for the tensor. Can be a scalar number, a nested Python list, or any DLPack-compatible array (NumPy, PyTorch, and so on). If not provided, exactly one of storage or state must be supplied.
  • dtype (DType | None) – The data type for the tensor elements. For DLPack arrays this defaults to the array’s own dtype; passing a conflicting value raises ValueError. For Python scalars and lists, defaults to DType.float32 on CPU and DType.bfloat16 on accelerators.
  • device (Device | DeviceMesh | None) – The device where the tensor will be allocated. Defaults to an accelerator if available, otherwise CPU. Only valid when data is provided.
  • storage (driver.Buffer | None) – Internal backing buffer for a realized tensor. Mutually exclusive with data.
  • state (RealizationState | None) – Internal realization state for an unrealized tensor. Mutually exclusive with data.

Return type:

Tensor

Allocates the tensor, delegating to F.constant when data is given.

When data is provided, returns the tensor produced by F.constant directly so that lazy/eager realization contexts track the correct object. For internal construction (storage or state), falls through to the normal allocation path.

T​

property T: Tensor

source

Returns a tensor with the last two dimensions transposed.

This is equivalent to calling transpose(-1, -2), which swaps the last two dimensions of the tensor. For a 2D matrix, this produces the standard matrix transpose.

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

# Create a 2x3 matrix
x = Tensor([[1, 2, 3], [4, 5, 6]], dtype=DType.int32)
print(f"Original shape: {x.shape}")
# Output: Original shape: [Dim(2), Dim(3)]

# Use .T property (equivalent to transpose(-1, -2))
y = x.T
print(f"Transposed shape: {y.shape}")
# Output: Transposed shape: [Dim(3), Dim(2)]
print(y)

Returns:

A tensor with the last two dimensions transposed.

arange()​

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

source

Creates a tensor with evenly spaced values within a given interval.

Returns a new 1-D tensor containing a sequence of values starting from start (inclusive) and ending before stop (exclusive), with values spaced by step.

Currently, graph compilation fails when stop - start isn’t evenly divisible by step. For example, Tensor.arange(0, 5, 2) should produce three values, [0, 2, 4], but shape inference declares an output length of 2. The generated values therefore don’t fit the declared output shape.

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

# Create a range from 0 to 10 (exclusive)
x = tensor.Tensor.arange(10)
# Result: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

# Create a range from 5 to 15 with step 2
y = tensor.Tensor.arange(5, 15, 2)
# Result: [5, 7, 9, 11, 13]

# Use a specific dtype
z = tensor.Tensor.arange(0, 5, dtype=DType.float32)
# Result: [0.0, 1.0, 2.0, 3.0, 4.0]

# Create a range with a floating-point step
w = tensor.Tensor.arange(0.0, 1.0, 0.25)
# Result: [0.0, 0.25, 0.5, 0.75]

# Create a descending range with negative step
v = tensor.Tensor.arange(5, 0, -1, dtype=DType.float32)
# Result: [5.0, 4.0, 3.0, 2.0, 1.0]

Parameters:

  • start (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The starting value of the sequence. If stop is not provided, this becomes the stop value and start defaults to 0.
  • stop (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – The end value of the sequence (exclusive). If not specified, the sequence ends at start and begins at 0.
  • step (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The spacing between values in the sequence. Must be non-zero.
  • out_dim (int | str | Dim | integer | TypedAttr | None) – The expected output dimension. Required when start, stop, or step are tensors rather than scalar literals. If not specified, the output dimension is computed from the scalar values of the inputs.
  • 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 | DeviceMesh | DeviceMapping | None) – The device where the tensor will be allocated. If not specified, defaults to an accelerator if available, otherwise CPU. Sharded placement is not supported.

Returns:

A 1-D Tensor containing the evenly spaced values.

Raises:

  • ValueError – If inputs aren’t scalar, or dynamic scalar inputs omit out_dim.
  • RuntimeError – If a statically known interval isn’t evenly divisible by step, causing the inferred output length to disagree with the number of generated values.

Return type:

Tensor

argmax()​

argmax(axis=-1)

source

Returns the indices of the maximum values along an axis.

It’s useful for finding the position of the largest element along a given dimension, such as determining predicted classes in classification.

When the input contains ties (identical maximum values), behavior depends on the device: CPU returns the first matching index, while GPU may return any of them.

from max.experimental import Tensor

x = Tensor([[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]])
indices = x.argmax(axis=-1)
# indices has shape (2, 1): [[1], [2]]

# Or flatten before reducing:
flat_index = x.argmax(axis=None)
# flat_index has shape (1,): [6] (flattened index of max value 4.2)

Parameters:

axis (int | None) – The axis along which to compute the argmax. Negative values index from the last dimension. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor with int64 dtype containing the indices of the maximum values along axis. For an integer axis, the result has the same rank as the input with the axis dimension reduced to size 1. When axis is None, the result has shape (1,).

Raises:

ValueError – If axis is out of range.

Return type:

Tensor

broadcast_to()​

broadcast_to(shape)

source

Broadcasts the tensor to a target shape.

Each input dimension must either equal the corresponding target dimension or be 1 (which is then stretched to match). This follows NumPy broadcasting semantics and is equivalent to PyTorch’s torch.broadcast_to() and torch.Tensor.expand().

from max.experimental import Tensor

x = Tensor.ones([3, 1])
result = x.broadcast_to([3, 4])
# result has shape (3, 4)

# Add a new leading dimension
result = x.broadcast_to([2, 3, 4])
# result has shape (2, 3, 4)

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The target shape. A static shape (no dynamic dimensions).

Returns:

A Tensor with the same elements as self but with the target shape.

Return type:

Tensor

buffers​

property buffers: tuple[Buffer, ...]

source

The underlying per-shard driver buffers.

Returns one buffer for non-distributed tensors, N buffers for a distributed tensor with N shards.

Raises:

TypeError – If the tensor is unrealized (lazy/symbolic).

cast()​

cast(dtype)

source

Casts the tensor to a different data type.

Returns a new tensor with the same values but a different data type. This is useful for type conversions between different numeric types, such as converting float32 to int32 for indexing operations or float32 to bfloat16 for memory-efficient computations.

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

# Create a float32 tensor
x = tensor.Tensor([1.7, 2.3, 3.9], dtype=DType.float32)
print(x.dtype)  # DType.float32

# Cast to int32 (truncates decimal values)
y = x.cast(DType.int32)
print(y.dtype)  # DType.int32
# Values: [1, 2, 3]

Parameters:

dtype (DType) – The target data type for the tensor.

Returns:

A new tensor with the specified data type, or self if the tensor already has the target dtype.

Return type:

Tensor

clip()​

clip(*, min=None, max=None)

source

Clips values outside a range to the boundaries of the range.

from max.experimental import tensor

# Create a 2x4 tensor
x = tensor.Tensor(
    [[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]],
)

# Find max along last axis (within each row)
clipped_above = x.clip(max=3.)
# Result: [[1.2, 3., 2.1, 0.8], [2.3, 1.9, 3, 3.]]

clipped_below = x.clip(min=3.)
# Result: [[3., 3.5, 3., 3.], [3., 3., 4.2, 3.]]

Parameters:

Returns:

A tensor containing the values clipped to the specified range.

Return type:

Tensor

constant()​

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

source

Creates a constant tensor from a Python literal or array-like value.

Parameters:

  • value (DLPackArray | Sequence[float | number[Any] | Sequence[Number | NestedArray]] | float | number[Any]) – The value to embed. A Python scalar, a (nested) sequence of numbers, or an array-like object that supports DLPack, such as a NumPy array.
  • dtype (DType | None) – The constant tensor’s element type. For an array-like value, defaults to the array’s dtype. For a Python scalar or sequence, defaults to DType.float32 on CPU or DType.bfloat16 on accelerators.
  • device (Device | None) – The device where the tensor is allocated. Defaults to an accelerator if available, otherwise the CPU.

Returns:

A Tensor containing the constant, with the same shape as value. A scalar value produces a rank-0 tensor.

Raises:

  • TypeError – If dtype is a sub-byte type.
  • ValueError – If value is a nested sequence that isn’t rectangular, if an integer in value is out of range for dtype, or if dtype doesn’t match the dtype of an array-like value.

Return type:

Tensor

device​

property device: Device

source

Gets the device where the tensor is stored.

Returns the device (CPU or accelerator) where the tensor’s data is located. Raises for distributed tensors that span multiple devices.

Returns:

The device where the tensor is stored.

Return type:

Device

driver_tensor​

property driver_tensor: Buffer

source

A pointer to the underlying memory.

Raises if the tensor is unrealized or sharded.

dtype​

property dtype: DType

source

Gets the data type of the tensor elements.

Returns:

The data type of the tensor elements.

Return type:

DType

external_name​

property external_name: str | None

source

The name a weight loads under, which for a plain tensor is nothing.

Only a weight is a weight, and it answers with its own name. Every other tensor – a constant, a random draw, the result of an op, or anything computed from a weight – is not one, which is what makes this the authoritative answer to “which registry entry loads here”.

Returns:

None.

from_dim()​

classmethod from_dim(dim)

source

Materializes a dimension as a rank-0 (scalar) tensor on CPU.

Converts a shape dimension — static, symbolic, or an algebraic expression such as batch * seq — into a scalar tensor holding its runtime value. This is the supported way to predicate runtime control flow on a symbolic dimension: a symbolic Dim cannot be compared to a Python int at trace time (int(dim) and dim <= 2 both fail for dynamic dims), but the materialized tensor can, and the comparison’s result is exactly the scalar boolean predicate that cond() expects.

from max.driver import CPU
from max.dtype import DType
from max.experimental import functional as F
from max.experimental.nn import Module
from max.experimental.sharding import TensorLayout
from max.experimental.tensor import Tensor
from max.graph import TensorType

class ScaleByBatch(Module):
    def forward(self, x: Tensor) -> Tensor:
        out_type = TensorType(x.dtype, x.shape, device=x.device)
        pred = Tensor.from_dim(x.shape[0]) <= 2  # scalar bool tensor on CPU
        then_fn, else_fn = lambda: x * 2.0, lambda: x * 4.0
        (out,) = F.cond(pred, [out_type], then_fn, else_fn)
        return out

model = ScaleByBatch().compile(
    TensorLayout(DType.float32, ["batch", 4], CPU())
)
result = model(Tensor.ones([1, 4], dtype=DType.float32, device=CPU()))

Parameters:

dim (int | str | Dim | integer | TypedAttr) – The dimension to materialize. Accepts anything DimLike (an int, a dim name, a Dim, or an algebraic dim expression).

Returns:

A rank-0 int64 tensor on CPU holding the dimension’s runtime value.

Return type:

Tensor

Raises:

ValueError – If the dimension has no value. This includes when the dimension is symbolic or algebraic in eager mode (no active graph), since it has no value outside a graph, or when it holds per-device sizes and no global size.

from_dlpack()​

classmethod from_dlpack(array)

source

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)

source

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 used for converting graph level values into tensor objects. The new tensor is registered as unrealized, backed by the current realization context.

Parameters:

value (Value[Any]) – 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

from_shard_values()​

classmethod from_shard_values(shard_values, mapping=None)

source

Creates a tensor from one value per device, moving no data.

Each value is one device’s shard: a graph value, or a single-device tensor in the graph being built. With one value and no mapping, the result is that value as a tensor. With several values, mapping is required and the result is a distributed tensor on the shards’ own devices.

Parameters:

Returns:

A tensor backed by the given shards.

Raises:

  • ValueError – If several shards are given without a mapping.
  • TypeError – If a shard is neither a graph value nor a tensor.

Return type:

Tensor

full()​

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

source

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 | TypedAttr]) – 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 | DeviceMesh | DeviceMapping | None) – The device, mesh, or device mapping where the tensor will be allocated. A mesh replicates the tensor. If not specified, defaults to the one set by default_device(), or else an accelerator if available, otherwise CPU. Pass a DeviceMapping to create a distributed tensor.

Returns:

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

Return type:

Tensor

full_like()​

classmethod full_like(input, value)

source

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

Returns a new tensor filled with the specified value that matches the shape, data type, and device of the input tensor. This behaves like NumPy’s full_like and PyTorch’s full_like.

from max.experimental import tensor

# Create a reference tensor
ref = tensor.Tensor.ones([2, 3])

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

Parameters:

  • input (Tensor) – The tensor to match. The returned tensor will have the same shape, dtype, and placement as it.
  • 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.

Return type:

Tensor

graph_values​

property graph_values: tuple[BufferValue | TensorValue, ...]

source

Returns per-shard graph values directly from the realization state.

For unrealized tensors (both distributed and single-device), returns the underlying graph values (TensorValue | BufferValue) without wrapping in intermediate Tensor objects.

For realized tensors, creates graph values via __tensorvalue__() on each shard.

Read this to reach the per-device graph values that custom dispatch rules and single-program multiple-data (SPMD) loops build on.

is_distributed​

property is_distributed: bool

source

Returns True when the tensor’s mesh holds more than one device.

That is a property of the mesh, not of the values: a fully replicated tensor on a four-device mesh reports True, and a tensor on a one-device mesh reports False whatever its placement.

item()​

item()

source

Gets the scalar value from a single-element tensor.

Extracts and returns the scalar value from a tensor containing exactly one element. The tensor is realized if needed and transferred to CPU before extracting the value.

For replicated distributed tensors, the value is read from the first shard (all shards hold identical data).

Returns:

The scalar value from the tensor. The return type matches the tensor’s dtype (such as float for float32, int for int32).

Raises:

  • TypeError – If the tensor contains more than one element.
  • ValueError – If the tensor is distributed and not fully replicated.

Return type:

Any

layout​

property layout: TensorLayout

source

This tensor’s dtype, global shape and distribution.

Immutable, so it says nothing about whether a callable may write to this tensor: that is declared where the tensor becomes an argument, by giving it a BufferLayout instead. Defined for a distributed tensor too, unlike type.

local_shards​

property local_shards: tuple[Tensor, ...]

source

Returns one unsharded tensor per device.

Each entry is a lightweight, unsharded Tensor backed by a single shard’s storage or graph value. They can be passed directly to F.* ops or used as Module parameters. The entries wrap the tensor’s own buffers rather than copies, so a write through one reaches the parent.

A realized distributed tensor gives one entry per driver.Buffer; an unrealized one gives one entry per GraphValue of the shared realization state. On a replicated tensor every entry holds the whole array, so read .shape on the entry for the per-device shape. A tensor on a single device returns (self,).

mapping​

property mapping: DeviceMapping

source

Returns the device mapping describing where this tensor lives.

materialize()​

materialize()

source

Gathers a distributed tensor into one local tensor on CPU.

Allreduces every Partial axis, allgathers every Sharded axis, then transfers the single copy to CPU. Returns self unchanged, on its own device, when the tensor is not distributed.

Raises:

ShardingError – If an Unknown placement would have to change.

Return type:

Tensor

max()​

max(axis=-1)

source

Computes the maximum along a specified axis.

from max.experimental import tensor

# Create a 2x4 tensor
x = tensor.Tensor(
    [[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]],
)

# Find max along last axis (within each row)
row_max = x.max(axis=-1)
# shape (2, 1): [[3.5], [4.2]]

# Find max along first axis (within each column)
col_max = x.max(axis=0)
# shape (1, 4): [[2.3, 3.5, 4.2, 3.1]]

# Find max over all elements
overall_max = x.max(axis=None)
# shape (1,): [4.2]

Parameters:

axis (int | None) – The axis along which to compute the maximum. Negative values index from the last dimension. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing the maximum along axis. For an integer axis, the result has the same rank as the input with the axis dimension reduced to size 1. When axis is None, the result has shape (1,).

Raises:

ValueError – If axis is out of range.

Return type:

Tensor

mean()​

mean(axis=-1)

source

Computes the mean along a specified axis.

from max.experimental import tensor

# Create a 2x4 tensor
x = tensor.Tensor(
    [[2.0, 4.0, 6.0, 8.0], [1.0, 3.0, 5.0, 7.0]],
)

# Compute mean along last axis (within each row)
row_mean = x.mean(axis=-1)
# shape (2, 1): [[5.0], [4.0]]

# Compute mean along first axis (within each column)
col_mean = x.mean(axis=0)
# shape (1, 4): [[1.5, 3.5, 5.5, 7.5]]

# Compute mean over all elements
overall_mean = x.mean(axis=None)
# shape (1,): [4.5]

Parameters:

axis (int | None) – The axis along which to compute the mean. Negative values index from the last dimension. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing the mean along axis. For an integer axis, the result has the same rank as the input with the axis dimension reduced to size 1; when axis is None, the result has shape (1,).

Raises:

ValueError – If axis is out of range.

Return type:

Tensor

mesh​

property mesh: DeviceMesh

source

Returns the device mesh.

min()​

min(axis=-1)

source

Computes the minimum along a specified axis.

from max.experimental import tensor

# Create a 2x4 tensor
x = tensor.Tensor(
    [[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]],
)

# Find min along last axis (within each row)
row_min = x.min(axis=-1)
# shape (2, 1): [[0.8], [1.9]]

# Find min along first axis (within each column)
col_min = x.min(axis=0)
# shape (1, 4): [[1.2, 1.9, 2.1, 0.8]]

# Find min over all elements
overall_min = x.min(axis=None)
# shape (1,): [0.8]

Parameters:

axis (int | None) – The axis along which to compute the minimum. Negative values index from the last dimension. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing the minimum along axis. For an integer axis, the result has the same rank as the input with the axis dimension reduced to size 1. When axis is None, the result has shape (1,).

Raises:

ValueError – If axis is out of range.

Return type:

Tensor

num_elements()​

num_elements()

source

Gets the total number of elements in the tensor.

Computes the product of all dimensions in the tensor’s shape to determine the total number of elements.

Returns:

The total number of elements in the tensor.

Return type:

int

num_shards​

property num_shards: int

source

Returns the number of per-device values.

That is the device count of the mesh, and 1 for a tensor on a single device. A replicated distributed tensor reports one value per device even though every value holds the whole array.

ones()​

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

source

Creates a tensor filled with ones.

Returns a new tensor with the specified shape where all elements are initialized to one.

from max.experimental import tensor

# Create a 2x3 tensor of ones
x = tensor.Tensor.ones((2, 3))

Parameters:

  • shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the output tensor.
  • 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 | DeviceMesh | DeviceMapping | None) – The device, mesh, or device mapping where the tensor will be allocated. A mesh replicates the tensor. If not specified, defaults to the one set by default_device(), or else an accelerator if available, otherwise CPU.

Returns:

A new tensor with the specified shape filled with ones.

Return type:

Tensor

ones_like()​

classmethod ones_like(input)

source

Creates a tensor of ones matching a given tensor’s properties.

Returns a new tensor filled with ones that matches the shape, data type, and device of the input tensor. This behaves like NumPy’s ones_like and PyTorch’s ones_like.

from max.experimental import tensor

# Create a reference tensor
ref = tensor.Tensor.zeros([3, 4])

# Create ones tensor matching the reference tensor
x = tensor.Tensor.ones_like(ref)
# Result: 3x4 tensor of ones with dtype float32

Parameters:

input (Tensor) – The tensor to match. The returned tensor will have the same shape, dtype, and placement as it.

Returns:

A new tensor filled with ones matching the properties of the input.

Return type:

Tensor

permute()​

permute(dims)

source

Permutes all dimensions of the tensor.

This is useful for changing the layout of multi-dimensional data, such as converting between different tensor layout conventions (for example, from [batch, channels, height, width] to [batch, height, width, channels]).

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

# x has shape (2, 3, 4): (batch, channels, length).
x = tensor.Tensor(
    [[[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]],
     [[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]]],
    dtype=DType.int32,
)
# Reorder to (batch, length, channels), producing shape (2, 4, 3).
y = x.permute([0, 2, 1])

Parameters:

dims (list[int]) – The target order of the dimensions as a list of axis indices. Each axis may be negative to index from the end of the tensor.

Returns:

A Tensor containing the input with its dimensions reordered to match dims. It has the same elements and dtype as the input, with the order of the elements changed according to the permutation.

Raises:

  • ValueError – If the length of dims does not match the rank of the input, or if dims contains duplicate dimensions.
  • IndexError – If any dimension in dims is out of range.

Return type:

Tensor

placements​

property placements: tuple[Placement, ...]

source

Returns one placement per mesh axis.

A NamedMapping carries the placements its spec resolved to when it was built, so this is a plain read. A spec that names one mesh axis for two tensor dimensions, or an axis that is both unreduced and sharded, raises ConversionError at construction.

prod()​

prod(axis=-1)

source

Computes the product along a specified axis.

Parameters:

axis (int | None) – The axis along which to compute the product. Negative values index from the last dimension. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing the product along axis. For an integer axis, the result has the same rank as the input with the axis dimension reduced to size 1. When axis is None, the result has shape (1,).

Raises:

ValueError – If axis is out of range.

Return type:

Tensor

range_like()​

classmethod range_like(input)

source

Creates a range tensor matching a given tensor’s properties.

Returns a new tensor containing sequential indices along the last dimension, broadcasted to match the shape of the given tensor. Each row (along the last dimension) contains values from 0 to the dimension size minus one. This is useful for creating position indices or coordinate tensors.

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

# Create a reference tensor with shape (2, 4)
ref = tensor.Tensor.zeros((2, 4), dtype=DType.int32, device=CPU())

# Create range tensor matching the reference tensor
x = tensor.Tensor.range_like(ref)
# Result: [[0, 1, 2, 3],
#          [0, 1, 2, 3]]

Parameters:

input (Tensor) – The tensor to match. The returned tensor will have the same shape, dtype, and placement as it, with values representing indices along the last dimension.

Returns:

A new tensor with sequential indices broadcasted to match the input’s shape.

Return type:

Tensor

rank​

property rank: int

source

Gets the number of dimensions in the tensor.

Returns the rank (number of dimensions) of the tensor. For example, a scalar has rank 0, a vector has rank 1, and a matrix has rank 2.

Returns:

The number of dimensions in the tensor.

Return type:

int

real​

property real: bool

source

Returns True if this tensor is realized (has concrete storage).

For sharded tensors this is all-or-nothing: either every shard is realized (_state is None) or none are.

realize​

property realize: Tensor

source

Force the tensor to realize if it is not already.

rebind_mapping()​

rebind_mapping(mapping)

source

Returns this tensor’s shards under a new placement, moving no data.

Use it to claim the placement of a per-device result, such as the Unknown output of an op without a sharding rule. The new placement is a trusted claim; ensure it is correct to prevent incorrect downstream results.

The mapping is checked against this tensor’s mesh and nothing else: a placement may name an axis the shards do not have, or a shape the shards do not match, and the tensor keeps the claim.

Parameters:

mapping (DeviceMapping) – The mapping to claim, on this tensor’s mesh.

Returns:

A tensor with the same shards and the given mapping.

Raises:

ValueError – If mapping is on a different mesh.

Return type:

Tensor

reshape()​

reshape(shape)

source

Reshapes the tensor.

If a value of -1 is present in shape, that dimension becomes an automatically calculated dimension collecting all unspecified dimensions. Its length becomes the number of elements in the original tensor divided by the product of the other dimensions of shape.

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

# x has shape (2, 3).
x = tensor.Tensor([[1, 2, 3], [4, 5, 6]], dtype=DType.int32)
# Flatten the same 6 elements into shape (6,).
y = x.reshape((6,))

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The new shape as an iterable of dimensions, such as a list or tuple of int or Dim values. A single dimension may be -1.

Returns:

A Tensor containing the input with a new shape. The order and total number of elements stays the same as the input.

Raises:

ValueError – If shape contains more than one -1 dimension, if a -1 dimension is requested while another dimension is 0, or if the input and target shapes have a different number of elements.

Return type:

Tensor

shape​

property shape: Shape

source

Gets the global (logical) shape of the tensor.

A distributed tensor folds its shape from the per-device shapes, and this works for unrealized tensors too. A sharded axis reports a per-shard dim that records the global extent, readable with global_dim; an Unknown axis whose per-device sizes differ keeps those sizes; and Replicated copies that disagree raise ValueError.

Returns:

The global shape of the tensor.

split()​

split(split_size_or_sections, axis=0)

source

Splits the tensor into multiple tensors along a given dimension.

This method supports two modes, matching PyTorch’s behavior:

  • If split_size_or_sections is an int, splits into chunks of that size (the last chunk may be smaller if not evenly divisible).
  • If split_size_or_sections is a list of ints, splits into chunks with exactly those sizes (must sum to the dimension size).
from max.experimental import tensor

# Create a 10x4 tensor
x = tensor.Tensor.ones([10, 4])

# Split into chunks of size 3 (last chunk is size 1)
chunks = x.split(3, axis=0)
# Result: 4 tensors with shapes [3,4], [3,4], [3,4], [1,4]

# Split into exact sizes
chunks = x.split([2, 3, 5], axis=0)
# Result: 3 tensors with shapes [2,4], [3,4], [5,4]

Parameters:

  • split_size_or_sections (int | list[int]) – Either an int (chunk size) or a list of ints (exact sizes for each output tensor).
  • axis (int) – The dimension along which to split. Defaults to 0.

Returns:

A list of Tensor objects resulting from the split.

Raises:

  • TypeError – If an integer chunk size is used for a non-static axis.
  • ValueError – If an explicit section size is negative or the section sizes don’t sum to the input size.
  • IndexError – If axis is out of range.

Return type:

list[Tensor]

squeeze()​

squeeze(axis)

source

Removes a dimension of size 1 from the tensor.

This is useful for removing singleton dimensions from tensors after operations that may have added them.

from max.experimental import tensor

# x has shape (4, 1, 6).
x = tensor.Tensor.ones([4, 1, 6])
# Remove the size-1 dimension at axis 1, producing shape (4, 6).
y = x.squeeze(axis=1)

Parameters:

axis (int) – The dimension to remove from the input’s shape. If negative, this indexes from the end of the tensor. For example, a value of -1 removes the last dimension.

Returns:

A Tensor containing the input with the dimension at axis removed. That dimension size must equal 1, so the result holds the same elements as the input with one fewer dimension.

Raises:

  • ValueError – If the dimension at axis does not have size 1.
  • IndexError – If axis is out of range, including for a rank-zero tensor.

Return type:

Tensor

state​

property state: RealizationState | None

source

Returns the realization state (unsharded tensors only).

Raises:

ValueError – If the tensor is distributed.

storage​

property storage: Buffer | None

source

Returns the single backing buffer (unsharded tensors only).

Raises:

ValueError – If the tensor is distributed.

sum()​

sum(axis=-1)

source

Computes the sum along a specified axis.

from max.experimental import tensor

# Create a 2x3 tensor
x = tensor.Tensor(
    [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
)

# Sum along last axis (within each row)
row_sum = x.sum(axis=-1)
# shape (2, 1): [[6.0], [15.0]]

# Sum along first axis (within each column)
col_sum = x.sum(axis=0)
# shape (1, 3): [[5.0, 7.0, 9.0]]

# Sum over all elements
total = x.sum(axis=None)
# shape (1,): [21.0]

Parameters:

axis (int | None) – The axis along which to compute the sum. Negative values index from the last dimension. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing the sum along axis. For an integer axis, the result has the same rank as the input with the axis dimension reduced to size 1. When axis is None, the result has shape (1,).

Raises:

ValueError – If axis is out of range.

Return type:

Tensor

to()​

to(target)

source

Transfers the tensor to a different device, mesh, or mapping.

This method supports three target types:

  1. Device: Transfers a single-device tensor to the target device. A realized tensor moves with to(); an unrealized one gets a transfer_to() op. A distributed tensor is allreduced and allgathered into a full copy first, then moved.
  2. DeviceMapping: Reassigns the tensor’s device mesh and placements. For single-device mappings, equivalent to .to(device). For multi-device mappings on a tensor on one device, the host slices each device’s part and transfers it, with no collective.
  3. DeviceMesh: Replaces the device mesh. Placements are kept, except that a tensor whose mapping is a NamedMapping re-resolves its named spec against the new mesh, so names the new mesh lacks become Replicated. For unsharded tensors targeting a multi-device mesh, creates a fully replicated mapping.
from max.experimental import tensor
from max.experimental.tensor import defaults
from max.driver import CPU

x = tensor.Tensor.ones((2, 3), device=CPU())
print(x.device)

_, mesh = defaults()
y = x.to(mesh.devices[0])
print(y.device)

z = y.to(y.device)
assert z is y

Parameters:

target (Device | DeviceMesh | DeviceMapping) –

The target for the tensor. Can be:

  • Device: Target device for transfer.
  • DeviceMesh: New mesh, keeping existing placements (or fully replicated for unsharded tensors).
  • DeviceMapping: New mesh and placements.

Returns:

A tensor on the specified target. Returns self when a realized tensor is already on the target device; an unrealized tensor on that device comes back as a new tensor with a transfer op.

Raises:

Return type:

Tensor

to_numpy()​

to_numpy()

source

Converts this tensor to a NumPy array.

Gathers a distributed tensor through the same allreduce and allgather path as materialize(), and moves a single-device tensor to CPU when it is stored elsewhere.

Raises:

ShardingError – If an Unknown placement would have to change.

Return type:

ndarray[Any, Any]

transpose()​

transpose(dim1, dim2)

source

Transposes two axes of the tensor.

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

# x has shape (2, 3).
x = tensor.Tensor([[1, 2, 3], [4, 5, 6]], dtype=DType.int32)
# Swap axes 0 and 1, producing shape (3, 2).
y = x.transpose(0, 1)

Parameters:

  • dim1 (int) – One of the two axes to transpose. If negative, this indexes from the end of the tensor. For example, a value of -1 refers to the last axis.
  • dim2 (int) – The other axis to transpose. If negative, this indexes from the end of the tensor.

Returns:

A Tensor containing the input with dim1 and dim2 transposed. It has the same elements and dtype as the input, with the order of the elements changed according to the transposition. For a rank-zero tensor, axes -1 and 0 are accepted and the scalar is returned unchanged.

Raises:

IndexError – If dim1 or dim2 is out of range.

Return type:

Tensor

type​

property type: TensorType

source

Gets the tensor type information.

Returns:

The type information for the tensor.

Return type:

TensorType

Raises:

TypeError – If the tensor is distributed.

unsqueeze()​

unsqueeze(axis)

source

Inserts a dimension of size 1 into the tensor.

This is the inverse of squeeze() and is useful for adding dimensions needed for broadcasting or matrix operations.

from max.experimental import tensor

# x has shape (3,).
x = tensor.Tensor([1.0, 2.0, 3.0])
# Add a size-1 dimension at the end, producing shape (3, 1).
y = x.unsqueeze(axis=-1)
# Add a size-1 dimension at the front, producing shape (1, 3).
z = x.unsqueeze(axis=0)

Parameters:

axis (int) – The index at which to insert a new dimension into the input’s shape. Elements at that index or higher are shifted back. If negative, it indexes relative to 1 plus the rank of the tensor. For example, a value of -1 adds a new dimension at the end, and -2 inserts the dimension immediately before the last dimension.

Returns:

A Tensor containing the input with a new dimension inserted at axis. That dimension has a size of 1, so the result holds the same elements as the input with one more dimension.

Raises:

ValueError – If axis is out of bounds.

Return type:

Tensor

zeros()​

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

source

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

# Create a 2x3 tensor of zeros
x = tensor.Tensor.zeros((2, 3))
# 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 | TypedAttr]) – 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 | DeviceMesh | DeviceMapping | None) – The device, mesh, or device mapping where the tensor will be allocated. A mesh replicates the tensor. If not specified, defaults to the one set by default_device(), or else 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(input)

source

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

Returns a new tensor filled with zeros that matches the shape, data type, and device of the input tensor. This behaves like NumPy’s zeros_like and PyTorch’s zeros_like.

from max.experimental import tensor

# Create a reference tensor
ref = tensor.Tensor.ones([3, 4])

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

Parameters:

input (Tensor) – The tensor to match. The returned tensor will have the same shape, dtype, and placement as it.

Returns:

A new tensor filled with zeros matching the properties of the input.

Return type:

Tensor

Models

Cookbook

Puzzles

Blog

Community