max.experimental.functional

`allgather()`

max.experimental.functional.allgather(t, tensor_axis=0, mesh_axis=0)

source

All-gathers a tensor’s shards along a mesh axis.

Transitions the tensor’s placement on mesh_axis from Sharded to Replicated. Each device gathers the shards from its peers and concatenates them along tensor_axis.

Parameters:

t (Tensor) – The input distributed tensor.
tensor_axis (int) – The tensor axis along which the shards are concatenated.
mesh_axis (int) – The mesh axis whose placement changes from Sharded to Replicated.

Returns:

A tensor with the full data replicated across mesh_axis.

Return type:

`allreduce_sum()`

max.experimental.functional.allreduce_sum(t, mesh_axis=0)

source

All-reduces a tensor by summing its shards across a mesh axis.

Transitions the tensor’s placement on mesh_axis from Partial to Replicated. Every device on mesh_axis ends up holding the sum of all inputs along that axis.

Parameters:

t (Tensor) – The input distributed tensor.
mesh_axis (int) – The mesh axis along which to reduce.

Returns:

A tensor with the same per-device values everywhere along mesh_axis.

Return type:

`any_distributed()`

max.experimental.functional.any_distributed(args)

source

True if any Tensor in args is distributed (multi-device).

Parameters:: args (tuple[object, ...])
Return type:: bool

`arange()`

max.experimental.functional.arange(start, stop, step=1, out_dim=None, *, dtype=None, device=None)

source

Creates a sequence of evenly spaced values from start to stop.

The sequence begins at start and increments by step, stopping before stop (the upper bound is exclusive).

stop - start must be zero or have the same sign as step. Currently, graph compilation fails when stop - start isn’t evenly divisible by step. For example, range(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.dtype import DType
from max.experimental import functional as F

result = F.range(0, 5, 1, dtype=DType.float32)
# result holds [0.0, 1.0, 2.0, 3.0, 4.0].

Parameters:

start (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The first value in the sequence. Must be a scalar value.
stop (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The exclusive upper bound. The sequence stops before this value. Must be a scalar value.
step (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The spacing between consecutive values. Must be non-zero. Defaults to 1.
out_dim (int | str | Dim | integer | TypedAttr | None) – The expected length of the output. Required when dynamic scalar tensor inputs prevent static length inference. When omitted, it’s computed from scalar literals.
dtype (DType | None) – The element type of the result tensor. Defaults to float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping. Sharded placement is not supported.

Returns:

A Tensor containing the generated sequence.

Raises:

ValueError – If out_dim is omitted for dynamic scalar inputs, if any input isn’t scalar, if any input isn’t on the CPU, or if device requests a sharded placement.
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:

`argmax()`

max.experimental.functional.argmax(x, 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
from max.experimental import functional as F

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

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

Parameters:

x (Tensor) – The input tensor.
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 x 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 for x.

Return type:

`argmin()`

max.experimental.functional.argmin(x, axis=-1)

source

Returns the indices of the minimum values along an axis.

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

from max.experimental import Tensor
from max.experimental import functional as F

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

Parameters:

x (Tensor) – The input tensor.
axis (int | None) – The axis along which to compute the argmin. 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 minimum values along axis. For an integer axis, the result has the same rank as x 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 for x.

Return type:

`argsort()`

max.experimental.functional.argsort(x, ascending=True)

source

Returns the indices that would sort a rank-1 tensor.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([3.0, 1.0, 2.0])
# Ascending order visits 1, 2, 3, so the indices are [1, 2, 0].
result = F.argsort(x, ascending=True)
# result is [1, 2, 0]

Parameters:

x (Tensor) – The input tensor to sort. Must be rank 1.
ascending (bool) – Whether to sort in ascending order. If False, sorts in descending order. Defaults to True.

Returns:

A Tensor containing the sorting indices, with the same shape as x and int64 dtype.

Raises:

ValueError – If x is not rank 1.

Return type:

`as_interleaved_complex()`

max.experimental.functional.as_interleaved_complex(x)

source

Reshapes a real tensor of alternating (real, imag) values into complex form.

Pulls each adjacent (real, imag) pair in the last dimension out into a trailing pair of size 2.

Parameters:: x (Tensor) – A real tensor representing complex numbers as alternating pairs of (real, imag) values. The last dimension must have an even size.
Returns:: A tensor of shape (*x.shape[:-1], x.shape[-1] // 2, 2). All dimensions except the last are unchanged; the last dimension is halved, and a final dimension of size 2 is appended to hold the (real, imag) components.
Return type:: Tensor

`atanh()`

max.experimental.functional.atanh(x)

source

Computes the inverse hyperbolic tangent of a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([-0.5, 0.0, 0.5])
result = F.atanh(x)
# result is approximately [-0.549, 0.0, 0.549]

Parameters:: x (Tensor) – The input tensor, with values in the range (-1, 1). Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing the inverse hyperbolic tangent of each element of x.
Return type:: Tensor

`avg_pool2d()`

max.experimental.functional.avg_pool2d(input, kernel_size, stride=1, dilation=1, padding=0, ceil_mode=False, count_boundary=True)

source

Applies 2D average pooling.

Slides a window of size kernel_size over the spatial dimensions and replaces each window with its average value.

Parameters:

input (Tensor) – The input tensor in channels-last (NHWC) layout, (batch_size, height, width, channels).
kernel_size (tuple[DimLike, DimLike]) – A tuple (kernel_h, kernel_w) giving the height and width of the sliding window.
stride (int | tuple[int, int]) – The stride of the sliding window. Either a single int applied to both spatial dimensions, or a tuple (stride_h, stride_w). Defaults to 1.
dilation (int | tuple[int, int]) – The spacing between kernel elements. Either a single int applied to both spatial dimensions, or a tuple (dilation_h, dilation_w). Defaults to 1.
padding (int | tuple[int, int]) – Zero-padding added to both sides of each spatial dimension. Either a single int applied to both spatial dimensions, or a tuple (pad_h, pad_w). Defaults to 0.
ceil_mode (bool) – When True, uses ceil instead of floor when computing the output spatial shape. Defaults to False.
count_boundary (bool) – When True, includes padding elements in the divisor when computing the average. Defaults to True.

Returns:

A Tensor containing the averaged values, with shape (batch_size, height_out, width_out, channels).

Return type:

`band_part()`

max.experimental.functional.band_part(x, num_lower=None, num_upper=None, exclude=False)

source

Set all values to zero except a diagonal band of an input matrix.

All but the last two axes are treated as batches, and the last two axes define the matrices.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
# Keep the main diagonal and one sub-diagonal, producing
# [[1, 0, 0], [1, 1, 0], [0, 1, 1]].
result = F.band_part(x, num_lower=1, num_upper=0)
# result is [[1, 0, 0], [1, 1, 0], [0, 1, 1]]

Parameters:

x (Tensor) – The input tensor to mask.
num_lower (int | None) – The number of diagonal bands to include below the central diagonal. If None or -1, includes the entire lower triangle. Defaults to None.
num_upper (int | None) – The number of diagonal bands to include above the central diagonal. If None or -1, includes the entire upper triangle. Defaults to None.
exclude (bool) – Whether to invert the selection, zeroing out the elements in the band instead. Defaults to False.

Returns:

A Tensor containing x with the masked-out elements set to zero and the remaining elements copied from x. It has the same shape and dtype as x.

Raises:

ValueError – If the input tensor rank is less than 2, or if num_lower or num_upper are out of bounds for statically known dimensions.

Return type:

`bottom_k()`

max.experimental.functional.bottom_k(input, k, axis=-1)

source

Returns the k smallest values along an axis with their indices.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, 3.0, 2.0, 5.0, 4.0])
# The two smallest values are 1 and 2 at indices 0 and 2.
values, indices = F.bottom_k(x, k=2, axis=-1)
# values is [1, 2]
# indices is [0, 2]

Parameters:

input (Tensor) – The input tensor from which to select the bottom k.
k (int) – The number of values to select from input. Must be in the range [0, input.shape[axis]].
axis (int) – The axis along which to select the bottom k. Defaults to -1. On a GPU input, only the last axis is supported.

Returns:

A tuple of two Tensor objects. The first holds the bottom k values along axis in ascending order, and the second holds their int64 indices in input. Both tensors have the shape of input with the axis dimension reduced to size k.

Return type:

tuple[Tensor, Tensor]

`broadcast_to()`

max.experimental.functional.broadcast_to(x, shape)

source

Broadcasts a 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().

from max.experimental import Tensor
from max.experimental import functional as F

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

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

Parameters:

x (Tensor) – The input tensor. Must not contain any dynamic dimensions.
shape (Iterable[int | str | Dim | integer | TypedAttr]) – The target shape. A static shape (no dynamic dimensions).

Returns:

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

Return type:

`buffer_store()`

max.experimental.functional.buffer_store(destination, source)

source

Stores values from a tensor into a tensor buffer.

Parameters:

destination (Tensor) – The destination buffer tensor.
source (Tensor) – The source tensor whose values are written into destination.

Return type:

None

`buffer_store_slice()`

max.experimental.functional.buffer_store_slice(destination, source, indices)

source

Stores values into a slice of a tensor buffer.

Parameters:

destination (Tensor) – The destination buffer tensor.
source (Tensor) – The source tensor whose values are written into the slice.
indices (SliceIndices) – The slice specification within destination to write to.

Return type:

None

`cast()`

max.experimental.functional.cast(x, dtype)

source

Casts a tensor to a different data type.

Values may change when the source and target types can’t represent each other exactly. Float-to-integer casts truncate toward zero; float-to-float casts with lower precision round to the nearest representable value.

from max.dtype import DType
from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.7, -1.7, 2.5])  # float32 on CPU by default
result = F.cast(x, DType.int32)
# result has dtype int32 and values [1, -1, 2]

Parameters:

x (Tensor) – The input tensor.
dtype (DType) – The target data type.

Returns:

A tensor with the same shape but the new dtype.

Return type:

`ceil()`

max.experimental.functional.ceil(x)

source

Computes the ceiling of a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.5, 2.0, -1.5, -2.7])
result = F.ceil(x)
# result is [2.0, 2.0, -1.0, -2.0]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing each element of x rounded up toward positive infinity.
Return type:: Tensor

`chunk()`

max.experimental.functional.chunk(x, chunks, axis=0)

source

Splits a tensor into equal-sized chunks along an axis.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1, 2, 3, 4, 5, 6])

# Split into three equal chunks along axis 0
parts = F.chunk(x, 3, axis=0)
# parts[0] is [1, 2]
# parts[1] is [3, 4]
# parts[2] is [5, 6]

Parameters:

x (Tensor) – The tensor to chunk.
chunks (int) – The number of chunks. Must be positive and evenly divide the size of x along axis.
axis (int) – The axis to split along. Defaults to 0.

Returns:

A list of Tensor objects (chunks), each the same size along axis.

Raises:

ValueError – If chunks does not evenly divide the size of x along axis, or if x is a scalar and chunks is greater than 1.
IndexError – If axis is out of range for x.

Return type:

`clamp()`

max.experimental.functional.clamp(x, lower_bound, upper_bound)

source

Clamps tensor values to [lower_bound, upper_bound].

Parameters:

x (Tensor)
lower_bound (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray)
upper_bound (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray)

Return type:

`clip()`

max.experimental.functional.clip(x, lower_bound, upper_bound)

source

Clamps tensor values to [lower_bound, upper_bound].

Parameters:

x (Tensor)
lower_bound (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray)
upper_bound (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray)

Return type:

`complex_mul()`

max.experimental.functional.complex_mul(lhs, rhs)

source

Multiplies two complex-valued tensors element-wise.

Both inputs must use the interleaved complex representation (trailing dimension of size 2).

Parameters:

lhs (Tensor) – The left-hand side complex tensor.
rhs (Tensor) – The right-hand side complex tensor.

Returns:

A complex tensor with the broadcast shape containing element-wise products.

Return type:

`concat()`

max.experimental.functional.concat(original_vals, axis=0)

source

Concatenates tensors along an axis.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([[1, 2], [3, 4]])
b = Tensor([[5, 6], [7, 8]])

vertical = F.concat([a, b], axis=0)
# vertical has shape (4, 2):
# [[1, 2], [3, 4], [5, 6], [7, 8]]

horizontal = F.concat([a, b], axis=1)
# horizontal has shape (2, 4):
# [[1, 2, 5, 6], [3, 4, 7, 8]]

Parameters:

original_vals (Iterable[Tensor]) – The tensors to concatenate. They must have the same rank and size on every dimension except axis.
axis (int) – The axis to concatenate along. Negative values count from the end. Defaults to 0.

Returns:

A Tensor containing the concatenated inputs. Its size along axis is the sum of the inputs’ sizes and every other axis is unchanged.

Raises:

ValueError – If no tensors are provided, if the inputs don’t all have the same rank, if they differ in size on any dimension other than axis, or if they aren’t all on the same device.
IndexError – If axis is out of range.

Return type:

`cond()`

max.experimental.functional.cond(pred, out_types, then_fn, else_fn)

source

Conditionally executes one of two branches based on a boolean predicate.

Both branches must return the same number and types of values as specified by out_types. The predicate is evaluated at runtime to determine which branch executes. If pred lives on a non-CPU device, it is transferred to CPU automatically.

from max.driver import CPU
from max.dtype import DType
from max.experimental import Tensor
from max.experimental import functional as F
from max.graph import DeviceRef, TensorType

def then_fn():
    return Tensor([1.0, 2.0], dtype=DType.float32, device=CPU())

def else_fn():
    return Tensor([10.0, 20.0], dtype=DType.float32, device=CPU())

pred = Tensor(True, dtype=DType.bool, device=CPU())
out_types = [TensorType(DType.float32, [2], DeviceRef.CPU())]
(result,) = F.cond(pred, out_types, then_fn, else_fn)
# pred is True, so result is [1.0, 2.0]

Parameters:

pred (Tensor) – A boolean scalar tensor of type bool determining which branch to execute.
out_types (Iterable[Type[Any]] | None) – The expected output types for both branches. Use None for branches that don’t return values (such as buffer mutations).
then_fn (Callable[[], Iterable[Tensor] | Tensor | None]) – A callable executed when pred is True.
else_fn (Callable[[], Iterable[Tensor] | Tensor | None]) – A callable executed when pred is False.

Returns:

The output values from the executed branch, or an empty list when out_types is None.

Return type:

`constant()`

max.experimental.functional.constant(value, dtype=None, device=None)

source

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

For an array-like value, the array’s own dtype is preserved when dtype is None. For a Python scalar or sequence, dtype defaults to float32 on CPU or bfloat16 on accelerators.

from max.dtype import DType
from max.experimental import functional as F

result = F.constant([[1.0, 2.0], [3.0, 4.0]], DType.float32)
# result holds [[1.0, 2.0], [3.0, 4.0]].

Warning

Loading a constant can lose precision. For example, loading 16777217 as a float32 produces 16777216.0.

Warning

When device is a multi-device DeviceMapping, the full value is materialized on every device. A mapping that shards a tensor axis therefore duplicates the data and reports a global shape scaled by the number of devices, without raising. Create the tensor Replicated and shard it afterward.

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 float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement. 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:

`constant_external()`

max.experimental.functional.constant_external(name, type, device=None, align=None)

source

Creates a constant tensor from external (weight) data.

External constants are loaded at compile time from the named weight rather than being inlined into the graph.

Parameters:

name (str) – The external symbol name to load (typically a weight identifier).
type (TensorType) – The TensorType describing the constant’s shape and dtype.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement.
align (int | None) – The alignment of the constant. If not provided, the default alignment for the type’s dtype will be used.

Returns:

A tensor on the requested placement initialized from the external data.

Return type:

`conv2d()`

max.experimental.functional.conv2d(x, filter, stride=(1, 1), dilation=(1, 1), padding=(0, 0, 0, 0), groups=1, bias=None, input_layout=ConvInputLayout.NHWC, filter_layout=FilterLayout.RSCF)

source

Computes the 2-D convolution product of the input with the given filter, bias, strides, dilations, paddings, and groups.

This uses the following layout assumptions:

The input has channels-last (NHWC) layout, meaning (batch_size, height, width, in_channels).
The filter has RSCF layout, meaning (height, width, in_channels / num_groups, out_channels).
The bias has shape (out_channels,).

The padding values are expected to take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after…) and represent padding 0’s before and after the indicated spatial dimensions in the input. In 2-D convolution, dim1 here represents H and dim2 represents W. In Python-like syntax, padding a 2x3 spatial input with [0, 1, 2, 1] would yield:

input = [
  [1, 2, 3],
  [4, 5, 6]
]
# Shape is 2x3

padded_input = [
  [0, 0, 1, 2, 3, 0],
  [0, 0, 4, 5, 6, 0],
  [0, 0, 0, 0, 0, 0]
]
# Shape is 3x6

This op currently only supports strides and padding on the input.

Convolving a 2x2 input with an all-ones 2x2 filter sums the window:

from max.experimental import Tensor
from max.experimental import functional as F

# NHWC input: batch 1, 2x2 spatial, 1 channel.
x = Tensor([[[[1.0], [2.0]], [[3.0], [4.0]]]])
# RSCF filter: 2x2, 1 in-channel, 1 out-channel, all ones.
filter = Tensor([[[[1.0]], [[1.0]]], [[[1.0]], [[1.0]]]])
result = F.conv2d(x, filter)
# result is [[[[10]]]] with shape (1, 1, 1, 1)

Parameters:

x (Tensor) – An NHWC input tensor to perform the convolution upon.
filter (Tensor) – The convolution filter in RSCF layout, (height, width, in_channels / num_groups, out_channels).
stride (tuple[int, int]) – The stride of the convolution operation.
dilation (tuple[int, int]) – The spacing between the kernel points.
padding (tuple[int, int, int, int]) – The amount of padding applied to the input.
groups (int) – When greater than 1, divides the convolution into multiple parallel convolutions. The number of input and output channels must both be divisible by the number of groups.
bias (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – An optional 1-D bias of shape (out_channels,).
input_layout (ConvInputLayout) – The layout of the input tensor. Defaults to NHWC.
filter_layout (FilterLayout) – The layout of the filter tensor. Defaults to RSCF.

Returns:

A Tensor containing the result of the convolution, with shape (batch_size, height_out, width_out, out_channels).

Raises:

ValueError – If x isn’t rank 4, filter isn’t rank 4, bias is given and isn’t rank 1, or x and filter aren’t on the same device.

Return type:

`conv2d_transpose()`

max.experimental.functional.conv2d_transpose(x, filter, stride=(1, 1), dilation=(1, 1), padding=(0, 0, 0, 0), output_paddings=(0, 0), bias=None, input_layout=ConvInputLayout.NHWC, filter_layout=FilterLayout.RSCF)

source

Computes the 2-D deconvolution of the input with the given filter, strides, dilations, and paddings.

This computes the transpose (gradient) of convolution, with the following layout assumptions (where out_channels is with respect to the original convolution):

The input x has channels-last (NHWC) layout, meaning (batch_size, height, width, in_channels).
The filter has RSCF layout, meaning (kernel_height, kernel_width, out_channels, in_channels).
The bias has shape (out_channels,).

This op effectively computes the gradient of a convolution with respect to its input, as if the original convolution had the same filter and hyperparameters as this op. For a visualization of the computation, see Transposed Convolution.

The padding values take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after, ...) and are cropped (removed) from the indicated spatial dimensions of the output. In 2-D transposed convolution, dim1 represents H_out and dim2 represents W_out. In Python-like syntax, cropping a 2x4 spatial output with [0, 1, 2, 1] would yield:

output = [
  [1, 2, 3, 4],
  [5, 6, 7, 8]
]
# Shape is 2x4

cropped_output = [
  [3],
]
# Shape is 1x1

Deconvolving a 1x1 input with an all-ones 2x2 filter (filter is RSCF, with out_channels and in_channels with respect to the original convolution):

from max.experimental import Tensor
from max.experimental import functional as F

# NHWC input: batch 1, 1x1 spatial, 1 channel.
x = Tensor([[[[3.0]]]])
# RSCF filter: 2x2 kernel, 1 out-channel, 1 in-channel, all ones.
filter = Tensor([[[[1.0]], [[1.0]]], [[[1.0]], [[1.0]]]])
result = F.conv2d_transpose(x, filter)

Parameters:

x (Tensor) – An NHWC input tensor to perform the deconvolution upon.
filter (Tensor) – The convolution filter in RSCF layout, (height, width, out_channels, in_channels).
stride (tuple[int, int]) – The stride of the sliding window as a tuple (stride_h, stride_w). Defaults to (1, 1).
dilation (tuple[int, int]) – The spacing between the kernel points.
padding (tuple[int, int, int, int]) – The amount cropped from each spatial dimension of the output.
output_paddings (tuple[int, int]) – The number of zeros added at the end of each output spatial axis. This resolves the ambiguity between multiple output shapes when a stride is greater than 1. Only 0 is supported.
bias (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – An optional tensor of shape (out_channels,).
input_layout (ConvInputLayout) – The layout of the input tensor. Defaults to NHWC.
filter_layout (FilterLayout) – The layout of the filter tensor. Defaults to RSCF.

Returns:

A Tensor containing the result of the deconvolution, in channels-first (NCHW) layout (batch_size, out_channels, height_out, width_out). This differs from the channels-last (NHWC) input layout.

Raises:

ValueError – If x isn’t rank 4, filter isn’t rank 4, bias is given and isn’t rank 1, an output padding isn’t smaller than its stride, or x and filter aren’t on the same device.

Return type:

`conv3d()`

max.experimental.functional.conv3d(x, filter, stride=(1, 1, 1), dilation=(1, 1, 1), padding=(0, 0, 0, 0, 0, 0), groups=1, bias=None, input_layout=ConvInputLayout.NHWC, filter_layout=FilterLayout.QRSCF)

source

Computes the 3-D convolution product of the input with the given filter, bias, strides, dilations, paddings, and groups.

This uses the following layout assumptions:

The input has channels-last (NDHWC) layout, meaning (batch_size, depth, height, width, in_channels).
The filter has QRSCF layout, meaning (depth, height, width, in_channels / num_groups, out_channels).

The padding values are expected to take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after…) and represent padding 0’s before and after the indicated spatial dimensions in the input. In 3-D convolution, dim1 here represents D, dim2 represents H and dim3 represents W. In Python-like syntax, padding a 2x3 spatial input with [0, 1, 2, 1] would yield:

input = [
  [1, 2, 3],
  [4, 5, 6]
]
# Shape is 2x3

padded_input = [
  [0, 0, 1, 2, 3, 0],
  [0, 0, 4, 5, 6, 0],
  [0, 0, 0, 0, 0, 0]
]
# Shape is 3x6

This op currently only supports strides and padding on the input.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

# NDHWC input: batch 1, 4x4x4 spatial, 1 channel.
x = Tensor.ones((1, 4, 4, 4, 1), dtype=DType.float32)
# QRSCF filter: 2x2x2, 1 in-channel, 1 out-channel.
filter = Tensor.ones((2, 2, 2, 1, 1), dtype=DType.float32)
result = F.conv3d(x, filter)
# result has shape (1, 3, 3, 3, 1)

Parameters:

x (Tensor) – An NDHWC input tensor to perform the convolution upon.
filter (Tensor) – The convolution filter in QRSCF layout, (depth, height, width, in_channels / num_groups, out_channels).
stride (tuple[int, int, int]) – The stride of the convolution operation.
dilation (tuple[int, int, int]) – The spacing between the kernel points.
padding (tuple[int, int, int, int, int, int]) – The amount of padding applied to the input.
groups (int) – When greater than 1, divides the convolution into multiple parallel convolutions. The number of input and output channels must both be divisible by the number of groups.
bias (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – An optional 1-D bias of shape (out_channels,).
input_layout (ConvInputLayout) – The layout of the input tensor. Defaults to NDHWC.
filter_layout (FilterLayout) – The layout of the filter tensor. Defaults to QRSCF.

Returns:

A Tensor containing the result of the convolution, with shape (batch_size, depth_out, height_out, width_out, out_channels).

Raises:

ValueError – If x isn’t rank 5, filter isn’t rank 5, or bias is given and isn’t rank 1.

Return type:

`cos()`

max.experimental.functional.cos(x)

source

Computes the cosine of a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([0.0, 0.5, 1.0])
result = F.cos(x)
# result is approximately [1.0, 0.878, 0.540]

Parameters:: x (Tensor) – The input interpreted as radians. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing the cosine of each element of x.
Return type:: Tensor

`cumsum()`

max.experimental.functional.cumsum(x, axis=-1, exclusive=False, reverse=False)

source

Computes the cumulative sum of a tensor along an axis.

Parameters:

x (Tensor) – The input tensor.
axis (int) – The axis along which to compute the cumulative sum. Defaults to -1.
exclusive (bool) – When True, the first output value is 0 and the final input element is excluded from the sum. Defaults to False.
reverse (bool) – When True, computes the sum starting from the end of the axis. Defaults to False.

Returns:

A tensor of the same shape and dtype where each element is the sum of the corresponding input elements up to that position along axis.

Return type:

`custom()`

max.experimental.functional.custom(name, device, values, out_types, parameters=None, custom_extensions=None)

source

Calls a custom op, optionally loading custom Mojo extensions first.

Parameters:

name (str) – The registered name of the custom op.
device (Device | DeviceRef) – The device on which to execute the op.
values (Sequence[Any]) – The input values passed to the op.
out_types (Sequence[Type[Any]]) – The expected output types.
parameters (Mapping[str, bool | int | str | DType] | None) – Optional compile-time parameters for the op.
custom_extensions (str | Path | Sequence[str | Path] | None) – Optional path or sequence of paths to custom Mojo extensions (.mojoc or .mojo sources) to load before invoking the op.

Returns:

A list of tensors produced by the custom op.

Return type:

`dequantize()`

max.experimental.functional.dequantize(encoding, quantized)

source

Dequantizes a quantized tensor to floating point.

Note

This currently supports the Q4_0, Q4_K, and Q6_K encodings only.

Parameters:

encoding (QuantizationEncoding) – The quantization encoding to use.
quantized (Tensor) – The quantized tensor to dequantize.

Returns:

A Tensor containing the dequantized, floating point result.

Raises:

ValueError – If encoding is not a supported quantization encoding, or if the last dimension isn’t divisible by the encoding’s block size.
TypeError – If the last dimension of quantized isn’t static.

Return type:

`distributed_broadcast()`

max.experimental.functional.distributed_broadcast(t, mesh)

source

Replicates a non-distributed tensor onto every device with one collective.

Parameters:

t (Tensor) – A non-distributed source tensor.
mesh (DeviceMesh) – The device mesh to replicate onto.

Returns:

A distributed tensor with Replicated placement on every axis.

Return type:

`div()`

max.experimental.functional.div(lhs, rhs)

source

Divides two tensors element-wise using true division (Python /).

For integer operands, this performs true division by promoting to float, matching Python’s / operator behavior. For floating-point operands, this performs standard floating-point division.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([10.0, 6.0, 3.0])
b = Tensor([2.0, 3.0, 4.0])
result = F.div(a, b)
# result is [5.0, 2.0, 0.75]

Parameters:

lhs (Tensor | int | float) – The numerator tensor or scalar.
rhs (Tensor | int | float) – The denominator tensor or scalar.

Returns:

A Tensor with the broadcast shape containing lhs / rhs element-wise. The result has a floating-point dtype for integer operands and the promoted dtype for mixed types.

Return type:

`elementwise_max()`

max.experimental.functional.elementwise_max(lhs, rhs)

source

Computes the element-wise maximum of two tensors.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([3.0, 1.0, 4.0])
b = Tensor([1.0, 5.0, 9.0])
result = F.elementwise_max(a, b)
# result is [3.0, 5.0, 9.0]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A tensor with the broadcast shape containing the larger value at each position.

Return type:

`elementwise_min()`

max.experimental.functional.elementwise_min(lhs, rhs)

source

Computes the element-wise minimum of two tensors.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([3.0, 1.0, 4.0])
b = Tensor([1.0, 5.0, 9.0])
result = F.elementwise_min(a, b)
# result is [1.0, 1.0, 4.0]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A tensor with the broadcast shape containing the smaller value at each position.

Return type:

`ensure_context()`

max.experimental.functional.ensure_context()

source

Ensures a realization context exists for Tensor / TensorValue conversion.

Return type:: Generator[None]

`equal()`

max.experimental.functional.equal(lhs, rhs)

source

Tests element-wise equality between two tensors.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([1.0, 2.0, 3.0])
b = Tensor([1.0, 5.0, 3.0])
result = F.equal(a, b)
# result is [True, False, True]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A Tensor with bool dtype containing the element-wise result of lhs == rhs.

Return type:

`erf()`

max.experimental.functional.erf(x)

source

Computes the error function of a tensor element-wise.

The error function erf is the probability that a randomly sampled normal distribution falls within a given range.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([-1.0, 0.0, 1.0])
result = F.erf(x)
# result is approximately [-0.843, 0.0, 0.843]

Parameters:: x (Tensor) – The input to the error function. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing the error function applied to each element of x.
Return type:: Tensor

`exp()`

max.experimental.functional.exp(x)

source

Computes the exponential of a tensor element-wise.

This applies exp(x) = e^x, where e is Euler’s number.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([0.0, 1.0, 2.0])
result = F.exp(x)
# result is approximately [1.0, 2.718, 7.389]

Parameters:: x (Tensor) – The input to the exponential function. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing e raised to the power of each element of x.
Return type:: Tensor

`flatten()`

max.experimental.functional.flatten(x, start_dim=0, end_dim=-1)

source

Flattens the specified dimensions of a tensor.

This does not change the order or total number of elements in the tensor. All dimensions from start_dim to end_dim (inclusive) are merged into a single output dimension.

from max.experimental import Tensor
from max.experimental import functional as F

# x has shape (2, 2, 2).
x = Tensor.ones([2, 2, 2])
# Merge dimensions 1 and 2 into one, producing shape (2, 4).
result = F.flatten(x, start_dim=1)

Parameters:

x (Tensor) – The input tensor to flatten.
start_dim (int) – The first dimension to flatten. Supports negative indexing. Defaults to 0.
end_dim (int) – The last dimension to flatten (inclusive). Supports negative indexing. Defaults to -1.

Returns:

A Tensor containing the start_dim through end_dim of x merged into one dimension.

Raises:

IndexError – If start_dim or end_dim is out of range.
ValueError – If start_dim comes after end_dim.

Return type:

`floor()`

max.experimental.functional.floor(x)

source

Computes the floor of a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.5, 2.0, -1.5, -2.7])
result = F.floor(x)
# result is [1.0, 2.0, -2.0, -3.0]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing each element of x rounded down toward negative infinity.
Return type:: Tensor

`floor_div()`

max.experimental.functional.floor_div(lhs, rhs)

source

Divides two tensors element-wise using floor division (Python //).

The result is rounded toward negative infinity, matching Python’s //. Either operand may be a Python int or float scalar, which is automatically promoted to a tensor. Integer operands stay in the integer domain (no float64 promotion), unlike div().

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

a = Tensor([7, 10, 18], dtype=DType.int32)
b = Tensor([2, 5, 6], dtype=DType.int32)
result = F.floor_div(a, b)
# result is [3, 2, 3]

# Floating-point operands are supported and still round toward -inf.
result = F.floor_div(Tensor([7.5, -7.5], dtype=DType.float32), 2.0)
# result is [3.0, -4.0]

Parameters:

lhs (Tensor | int | float) – The numerator tensor or scalar.
rhs (Tensor | int | float) – The denominator tensor or scalar.

Returns:

A Tensor with the broadcast shape containing the element-wise floor division of lhs by rhs.

Return type:

`fold()`

max.experimental.functional.fold(input, output_size, kernel_size, stride=1, dilation=1, padding=0)

source

Combines an array of sliding local blocks into a larger tensor.

L, the number of blocks, must equal prod((output_size[d] + 2 * padding[d] - dilation[d] * (kernel_size[d] - 1) - 1) // stride[d] + 1), where d ranges over all spatial dimensions.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

# Shape (N, C * kernel_h * kernel_w, L) = (1, 1 * 2 * 2, 9).
x = Tensor.ones((1, 4, 9), dtype=DType.float32)
# Fold nine 2x2 blocks into a 4x4 image.
result = F.fold(x, output_size=(4, 4), kernel_size=(2, 2))
# result has shape (1, 1, 4, 4)

Parameters:

input (Tensor) – The 3-D tensor to fold, with shape (N, C * kernel_sizes, L), where N is the batch dimension, C is the number of channels, kernel_sizes is the product of the kernel sizes, and L is the number of local blocks.
output_size (tuple[DimLike, DimLike]) – The spatial dimensions of the output tensor, as a tuple of two ints.
kernel_size (tuple[DimLike, DimLike]) – The size of the sliding blocks, as a tuple of two ints.
stride (int | tuple[int, int]) – The stride of the sliding blocks. Either an int or a tuple of two ints. Defaults to 1.
dilation (int | tuple[int, int]) – The spacing between kernel elements. Either an int or a tuple of two ints. Defaults to 1.
padding (int | tuple[int, int]) – The zero-padding added on both sides of the input. Either an int or a tuple of two ints. Defaults to 0.

Returns:

A Tensor containing the folded 4-D tensor, with shape (N, C, output_size[0], output_size[1]).

Raises:

ValueError – If the input’s channel dimension isn’t a multiple of the total kernel size, or if the number of blocks L doesn’t match the value computed from the other arguments.

Return type:

`full()`

max.experimental.functional.full(shape, value, *, dtype=None, device=None)

source

Creates a tensor filled with a single value.

When device is a DeviceMapping, the result is distributed across that mesh according to its placements.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the resulting tensor.
value (float | number[Any]) – The fill value.
dtype (DType | None) – The data type of the tensor. Defaults to float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement. Defaults to the current realization context’s device.

Returns:

A tensor of the requested shape, dtype, and placement with every element set to value.

Return type:

`full_like()`

max.experimental.functional.full_like(like, value)

source

Creates a tensor filled with a single value, matching another tensor’s shape and dtype.

Parameters:

like (Tensor | TensorType | DistributedTensorType) – The template tensor whose shape, dtype, and placement are copied.
value (float | number[Any]) – The fill value.

Returns:

A tensor matching the shape, dtype, and placement of like, with every element set to value.

Return type:

`functional()`

max.experimental.functional.functional(graph_op, rule=None)

source

Wraps a graph op as a distributed dispatch entry.

Returns a callable that local-auto-shards when any argument is a distributed Tensor (and a rule is bound), and otherwise forwards to the bare graph_op. The returned wrapper carries graph_op and rule as attributes; reassign wrapper.rule to swap the sharding rule at runtime without re-wrapping.

Parameters:

graph_op (Callable[[...], Any])
rule (Callable[[...], ActionSet] | None)

Return type:

Callable[[…], Any]

`gather()`

max.experimental.functional.gather(input, indices, axis)

source

Selects elements out of an input tensor by index.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

x = Tensor([[1, 2], [3, 4], [5, 6]], dtype=DType.int32)
indices = Tensor([0, 2], dtype=DType.int64)
# Select rows 0 and 2, producing [[1, 2], [5, 6]].
result = F.gather(x, indices, axis=0)
# result is [[1, 2], [5, 6]]

Note

When the gather axis is Sharded, the dispatcher first calls allgather() to make the input Replicated. It doesn’t emit an expert-parallel (Sharded(a_axis), R) → Partial(SUM) row, because that’s only correct when the caller masks indices per rank. Models that want expert-parallel semantics override gather.rule with their own rule.

Parameters:

input (Tensor) – The input tensor to select elements from.
indices (Tensor) – A tensor of int32 or int64 index values on the same device as input.
axis (int) – The dimension that indices indexes into input. If negative, indexes relative to the end of the input tensor. For example, gather(input, indices, axis=-1) indexes against the last dimension of input.

Returns:

A Tensor containing the selected elements. Its shape is input.shape with the dimension at axis replaced by indices.shape.

Raises:

IndexError – If axis is out of range for input.
ValueError – If indices isn’t integral or isn’t on the same device as input.

Return type:

`gather_nd()`

max.experimental.functional.gather_nd(input, indices, batch_dims=0)

source

Selects elements from a tensor by N-dimensional index.

Unlike gather(), which indexes along a single axis, gather_nd() indexes along multiple dimensions at once. The last dimension of indices is the index vector: its values select elements from input immediately after any batch_dims leading dimensions. Any remaining trailing dimensions of input are sliced into the output as features.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

input = Tensor(
    [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]
)
indices = Tensor([[0, 1], [1, 0]], dtype=DType.int64)
gathered = F.gather_nd(input, indices)
# gathered is [[3.0, 4.0], [5.0, 6.0]]

Each row of indices selects one row from the first two dimensions of input. The trailing dimension is copied into the output.

Parameters:

input (Tensor) – The input tensor to gather from.
indices (Tensor) – An integer tensor of multi-dimensional indices. Its last dimension must be static and gives the size of the index vector.
batch_dims (int) – The number of leading batch dimensions shared by input and indices. The shapes must match exactly along these leading dimensions. This function does not broadcast. Defaults to 0.

Returns:

A Tensor containing the gathered elements, with the same dtype as input. Its shape is the concatenation of:

input.shape[:batch_dims] — the leading batch dimensions.
indices.shape[batch_dims:-1] — the gather dimensions.
input.shape[batch_dims + indices.shape[-1]:] — the trailing sliced dimensions.

Raises:

ValueError – If any input is invalid. This includes when indices’s last dimension is not static, indices is not an integer tensor, batch_dims is negative or greater than indices.rank - 1, batch_dims + indices.shape[-1] exceeds input.rank, or the leading batch_dims of input and indices don’t match.

Return type:

`gaussian()`

max.experimental.functional.gaussian(shape=(), mean=0.0, std=1.0, *, dtype=None, device=None)

source

Samples values from a Gaussian (normal) distribution with the given mean and standard deviation.

When device is a DeviceMapping, each Sharded axis draws an independent stream while shards on Replicated axes draw identical values.

from max.experimental import functional as F

result = F.gaussian((2, 3), mean=0.0, std=1.0)
# result is a (2, 3) tensor sampled from a standard normal distribution.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the resulting tensor.
mean (float) – The mean of the distribution. Defaults to 0.0.
std (float) – The standard deviation of the distribution. Defaults to 1.0.
dtype (DType | None) – The data type of the tensor.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement.

Returns:

A Tensor of the requested shape, dtype, and placement with values sampled from Normal(mean, std**2).

Return type:

`gaussian_like()`

max.experimental.functional.gaussian_like(like, mean=0.0, std=1.0)

source

Samples Gaussian values matching another tensor’s shape and dtype.

Parameters:

like (Tensor | TensorType | DistributedTensorType) – The template tensor whose shape, dtype, and placement are copied.
mean (float) – The mean of the distribution. Defaults to 0.0.
std (float) – The standard deviation of the distribution. Defaults to 1.0.

Returns:

A tensor matching the shape, dtype, and placement of like, with values sampled from Normal(mean, std**2).

Return type:

`gelu()`

max.experimental.functional.gelu(x, approximate='none')

source

Applies the GELU (Gaussian Error Linear Unit) activation element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([-1.0, 0.0, 1.0])
result = F.gelu(x)
# result is approximately [-0.159, 0.0, 0.841]

Parameters:

x (Tensor) – The input to the GELU computation. Must have a floating-point dtype.
approximate (str) – The approximation method. Defaults to "none" (exact form using erf). Use "tanh" for the tanh-based approximation or "quick" for the sigmoid-based approximation.

Returns:

A Tensor of the same shape and dtype as x containing the GELU activation applied to each element of x.

Return type:

`greater()`

max.experimental.functional.greater(lhs, rhs)

source

Tests element-wise whether one tensor is greater than another.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([1.0, 5.0, 3.0])
b = Tensor([2.0, 3.0, 3.0])
result = F.greater(a, b)
# result is [False, True, False]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A Tensor with bool dtype containing the element-wise result of lhs > rhs.

Return type:

`greater_equal()`

max.experimental.functional.greater_equal(lhs, rhs)

source

Tests element-wise whether one tensor is greater than or equal to another.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([1.0, 5.0, 3.0])
b = Tensor([2.0, 3.0, 3.0])
result = F.greater_equal(a, b)
# result is [False, True, True]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A Tensor with bool dtype containing the element-wise result of lhs >= rhs.

Return type:

`group_norm()`

max.experimental.functional.group_norm(input, gamma, beta, num_groups, epsilon)

source

Computes group normalization over the channel axis of input.

Splits the channel axis (axis 1) of input into num_groups groups, computes the mean and variance within each group, and normalizes. gamma and beta then apply a per-channel affine transform. Useful when the batch axis is small enough that batch normalization is unstable.

Note

This op executes only on CUDA/HIP GPU targets.

Parameters:

input (Tensor) – The tensor to normalize, of shape (batch, channels, ...).
gamma (Tensor) – The per-channel scale applied after normalization. A 1-D tensor whose length matches the channel axis of input.
beta (Tensor) – The per-channel bias added after scaling. A 1-D tensor with the same shape as gamma.
num_groups (int) – The number of groups to split the channel axis into. Must divide the channel size evenly.
epsilon (float) – A small positive constant added to the variance for numerical stability.

Returns:

A Tensor with the same shape and dtype as input.

Raises:

ValueError – If input has fewer than 2 dimensions.

Return type:

`hann_window()`

max.experimental.functional.hann_window(window_length, *, periodic=True, dtype=None, device=None)

source

Computes a Hann window of a given length.

For a symmetric window of N points with N > 1, the value at index n is:

w[n] = 0.5 * (1 - cos(2 * pi * n / (N - 1)))

A window of length 0 is empty, and a window of length 1 is [1].

A periodic window instead computes N + 1 points and drops the last one, which makes it suitable for spectral analysis.

from max.experimental import functional as F

result = F.hann_window(4, periodic=True)
# result holds [0.0, 0.5, 1.0, 0.5].

Parameters:

window_length (int) – The number of points in the window.
periodic (bool) – Whether to return a periodic window. When True, computes a symmetric window of window_length + 1 points and drops the last point, so that hann_window(L, periodic=True) equals hann_window(L + 1, periodic=False)[:-1]. Defaults to True.
dtype (DType | None) – The output tensor’s data type. Defaults to float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping. Sharded placement is not supported.

Returns:

A 1-D Tensor of shape (window_length,) containing the window.

Raises:

ValueError – If window_length is negative, or if device requests a sharded placement.
TypeError – If window_length isn’t an integer.

Return type:

`in_graph_context()`

max.experimental.functional.in_graph_context()

source

Returns True when executing inside a Graph context.

Return type:: bool

`inplace_custom()`

max.experimental.functional.inplace_custom(name, device, values, out_types=None, parameters=None, custom_extensions=None)

source

Calls an in-place custom op that mutates one or more of its inputs.

Like custom(), but for ops that mutate buffer values rather than returning new tensors.

Parameters:

name (str) – The registered name of the custom op.
device (Device | DeviceRef) – The device on which to execute the op.
values (Sequence[Any]) – The input values; one or more are mutated in place.
out_types (Sequence[Type[Any]] | None) – Optional expected output types. Most in-place ops return no outputs and can leave this as None.
parameters (dict[str, bool | int | str | DType] | None) – Optional compile-time parameters for the op.
custom_extensions (str | Path | Sequence[str | Path] | None) – Optional path or sequence of paths to custom Mojo extensions to load before invoking the op.

Returns:

A list of tensors produced by the custom op, or an empty list when the op produces no outputs.

Return type:

`irfft()`

max.experimental.functional.irfft(input_tensor, n=None, axis=-1, normalization=Normalization.BACKWARD, input_is_complex=False, buffer_size_mb=512)

source

Computes the inverse of the real-input FFT.

Parameters:

input_tensor (Tensor) – The input tensor to compute the inverse real FFT of.
n (int | None) – The size of the output tensor. The input tensor is padded or truncated to n // 2 + 1 along axis.
axis (int) – The axis along which to compute the inverse FFT. Defaults to -1.
normalization (Normalization | str) – The normalization to apply to the output tensor. One of "backward", "ortho", or "forward". When "backward", the output is divided by n. When "ortho", the output is divided by sqrt(n). When "forward", no normalization is applied.
input_is_complex (bool) – Whether the input tensor is already interleaved complex. When True, the last dimension of the input tensor must be 2, and is excluded from the dimension referred to by axis.
buffer_size_mb (int) – The estimated size of a persistent buffer to use for storage of intermediate results. Needs to be the same across multiple calls to irfft within the same graph.

Returns:

A real tensor that is the inverse FFT of the complex input. The shape matches the input shape, except along axis, which is replaced by n.

`is_inf()`

max.experimental.functional.is_inf(x)

source

Tests element-wise whether a tensor contains infinite values.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, float("inf"), float("-inf"), float("nan")])
result = F.is_inf(x)
# result is [False, True, True, False]

Parameters:: x (Tensor) – The input tensor.
Returns:: A Tensor with bool dtype and the same shape as x that is True where x is positive or negative infinity.
Return type:: Tensor

`is_nan()`

max.experimental.functional.is_nan(x)

source

Tests element-wise whether a tensor contains NaN values.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, float("inf"), float("nan"), 0.0])
result = F.is_nan(x)
# result is [False, False, True, False]

Parameters:: x (Tensor) – The input tensor.
Returns:: A Tensor with bool dtype and the same shape as x that is True where x is NaN.
Return type:: Tensor

`layer_norm()`

max.experimental.functional.layer_norm(input, gamma, beta, epsilon)

source

Computes layer normalization over the last dimension of input.

The output is gamma * (input - mean) / sqrt(var + epsilon) + beta, where mean and var are reduced over the last axis of input and broadcast back across the leading axes.

Reduction is performed in the dtype of input. For numerically stable normalization on float16 or bfloat16 inputs, cast to float32 before calling this op and cast the result back.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1.0, 3.0]])
gamma = Tensor([1.0, 1.0])
beta = Tensor([0.0, 0.0])
result = F.layer_norm(x, gamma, beta, epsilon=1e-5)
# Each row is normalized to zero mean and approximately unit variance.

Parameters:

input (TensorValue) – The tensor to normalize. Reduction runs over the last axis.
gamma (Tensor) – The scale applied after normalization. A 1-D tensor whose length matches the last dimension of input.
beta (Tensor) – The bias added after scaling. A 1-D tensor with the same shape as gamma.
epsilon (float) – A small positive constant added to the variance for numerical stability.

Returns:

A Tensor with the same shape and dtype as input.

Raises:

ValueError – If gamma or beta does not match the last dimension of input, or if epsilon is not positive.

Return type:

`lazy()`

max.experimental.functional.lazy()

source

Defers tensor realization until explicitly awaited.

Return type:: Generator[None]

`log()`

max.experimental.functional.log(x)

source

Computes the natural logarithm of a tensor element-wise.

This applies log(x). It is the inverse of the exponential function x = e^y, where e is Euler’s number. Note that log(x) is undefined for x <= 0 and complex numbers are not currently supported.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, 2.718, 7.389, 20.0])
result = F.log(x)
# result is approximately [0.0, 1.0, 2.0, 2.996]

Parameters:: x (Tensor) – The input to the log computation. Must have a floating-point dtype and contain positive values only.
Returns:: A Tensor of the same shape and dtype as x containing the natural logarithm of each element of x.
Return type:: Tensor

`log1p()`

max.experimental.functional.log1p(x)

source

Computes log(1 + x) element-wise.

Note that log(1 + x) is undefined for x <= -1 and complex numbers are not currently supported.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([0.0, 1e-7, 1.0])
result = F.log1p(x)
# result is approximately [0.0, 1e-7, 0.693]

Parameters:: x (Tensor) – The input to the log computation. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing log(1 + x) for each element of x.
Return type:: Tensor

`logical_and()`

max.experimental.functional.logical_and(lhs, rhs)

source

Computes the element-wise logical AND of two boolean tensors.

Only supports boolean inputs.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

a = Tensor([True, True, False], dtype=DType.bool)
b = Tensor([True, False, False], dtype=DType.bool)
result = F.logical_and(a, b)
# result is [True, False, False]

Parameters:

lhs (Tensor | int | float) – The left-hand side boolean tensor.
rhs (Tensor | int | float) – The right-hand side boolean tensor.

Returns:

A Tensor with bool dtype containing the element-wise logical AND of lhs and rhs.

Return type:

`logical_not()`

max.experimental.functional.logical_not(x)

source

Computes the element-wise logical NOT of a boolean tensor.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

x = Tensor([True, False, True], dtype=DType.bool)
result = F.logical_not(x)
# result is [False, True, False]

Parameters:: x (Tensor) – The input boolean tensor.
Returns:: A Tensor with bool dtype and the same shape as x containing the element-wise logical NOT of x.
Return type:: Tensor

`logical_or()`

max.experimental.functional.logical_or(lhs, rhs)

source

Computes the element-wise logical OR of two boolean tensors.

Only supports boolean inputs.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

a = Tensor([True, True, False], dtype=DType.bool)
b = Tensor([True, False, False], dtype=DType.bool)
result = F.logical_or(a, b)
# result is [True, True, False]

Parameters:

lhs (Tensor | int | float) – The left-hand side boolean tensor.
rhs (Tensor | int | float) – The right-hand side boolean tensor.

Returns:

A Tensor with bool dtype containing the element-wise logical OR of lhs and rhs.

Return type:

`logical_xor()`

max.experimental.functional.logical_xor(lhs, rhs)

source

Computes the element-wise logical XOR of two boolean tensors.

Only supports boolean inputs.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

a = Tensor([True, True, False], dtype=DType.bool)
b = Tensor([True, False, False], dtype=DType.bool)
result = F.logical_xor(a, b)
# result is [False, True, False]

Parameters:

lhs (Tensor | int | float) – The left-hand side boolean tensor.
rhs (Tensor | int | float) – The right-hand side boolean tensor.

Returns:

A Tensor with bool dtype containing the element-wise logical XOR of lhs and rhs.

Return type:

`logsoftmax()`

max.experimental.functional.logsoftmax(value, axis=-1)

source

Computes the log-softmax of a tensor along an axis.

Parameters:

value (Tensor) – The input to the log-softmax computation. Must have a floating-point dtype.
axis (int) – The axis along which to compute the log-softmax. Defaults to the final axis (-1).

Returns:

A Tensor of the same shape and dtype as value containing the log-softmax of value computed along axis.

Return type:

`map_tensors()`

max.experimental.functional.map_tensors(fn, args)

source

Applies fn to every Tensor leaf in args.

Recurses into list and tuple containers; non-tensor leaves pass through unchanged.

Parameters:

fn (Callable[[Tensor], Any])
args (tuple[Any, ...])

Return type:

tuple[Any, …]

`masked_scatter()`

max.experimental.functional.masked_scatter(input, mask, updates, out_dim)

source

Updates tensor values at positions where mask is true.

Positions are filled in row-major order, so the first True position in mask takes the first element of updates, and so on.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

x = Tensor([[1, 2], [3, 4]], dtype=DType.int32)
mask = Tensor([[True, False], [False, True]], dtype=DType.bool)
updates = Tensor([10, 20], dtype=DType.int32)
# Write into the True positions, producing [[10, 2], [3, 20]].
result = F.masked_scatter(x, mask, updates, out_dim="num_updates")
# result is [[10, 2], [3, 20]]

Parameters:

input (Tensor) – The input tensor to write elements to.
mask (Tensor) – A tensor selecting the positions to write, broadcast to the shape of input. Pass a boolean tensor. A weak Python value is converted to boolean, but an existing tensor is used unchanged.
updates (Tensor) – A tensor of elements to write to input.
out_dim (int | str | Dim | integer | TypedAttr) – The new data-dependent dimension for the number of True positions in mask.

Returns:

A Tensor containing input with updates written where mask is true. It has the same shape and dtype as input.

Raises:

ValueError – If input and updates have mismatched dtypes, or if the inputs aren’t all on the same device.

Return type:

`matmul()`

max.experimental.functional.matmul(lhs, rhs)

source

Performs matrix multiplication between two tensors.

Treats the innermost two dimensions of each input as a matrix: lhs of shape (..., M, K) and rhs of shape (..., K, N) produce an output of shape (..., M, N). The K dimensions must match. Any outer batch dimensions are broadcast.

When inputs are distributed across devices, the operation is sharded according to the matmul sharding rule.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([[1.0, 2.0], [3.0, 4.0]])
b = Tensor([[5.0, 6.0], [7.0, 8.0]])
result = F.matmul(a, b)
# result has shape (2, 2):
# [[19.0, 22.0], [43.0, 50.0]]

# The ``@`` operator on Tensor also calls matmul.
result = a @ b

Parameters:

lhs (Tensor) – The left-hand side input tensor.
rhs (Tensor) – The right-hand side input tensor.

Returns:

A tensor representing the matrix product of lhs and rhs.

Return type:

`max()`

max.experimental.functional.max(x, y=None, /, axis=-1)

source

Computes the maximum of a tensor, or the element-wise maximum of two tensors.

Called with one argument, reduces x along axis. Called with two tensor arguments, returns their element-wise maximum.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]])

row_max = F.max(x, axis=-1)
# row_max has shape (2, 1): [[3.5], [4.2]]

col_max = F.max(x, axis=0)
# col_max has shape (1, 4): [[2.3, 3.5, 4.2, 3.1]]

y = Tensor([[2.0, 2.0, 2.0, 2.0], [2.0, 2.0, 2.0, 2.0]])
element_wise = F.max(x, y)
# element_wise: [[2.0, 3.5, 2.1, 2.0], [2.3, 2.0, 4.2, 3.1]]

Parameters:

x (Tensor) – The input tensor.
y (Tensor | None) – Optional second tensor. When provided, the result is the element-wise maximum of x and y.
axis (int | None) – The axis to reduce along when y is omitted. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing either the reduced maximum along axis or the element-wise maximum with the broadcast shape of the inputs.

Raises:

ValueError – If axis is out of range for x when reducing.

Return type:

`max_pool2d()`

max.experimental.functional.max_pool2d(input, kernel_size, stride=1, dilation=1, padding=0, ceil_mode=False)

source

Applies 2D max pooling to a tensor.

Slides a window of size kernel_size over the spatial dimensions and replaces each window with its maximum value.

Parameters:

input (Tensor) – The input tensor in channels-last (NHWC) layout, (batch_size, height, width, channels).
kernel_size (tuple[DimLike, DimLike]) – A tuple (kernel_h, kernel_w) giving the height and width of the sliding window.
stride (int | tuple[int, int]) – The stride of the sliding window. Either a single int applied to both spatial dimensions, or a tuple (stride_h, stride_w). Defaults to 1.
dilation (int | tuple[int, int]) – The spacing between kernel elements. Either a single int applied to both spatial dimensions, or a tuple (dilation_h, dilation_w). Defaults to 1.
padding (int | tuple[int, int]) – Padding added to both sides of each spatial dimension. Out-of-bounds positions are excluded from the maximum (equivalently, they use the dtype’s minimum value or negative infinity), so padding cannot win over negative input values. Either a single int applied to both spatial dimensions, or a tuple (pad_h, pad_w). Defaults to 0.
ceil_mode (bool) – When True, uses ceil instead of floor when computing the output spatial shape. Defaults to False.

Returns:

A Tensor containing the max-pooled values, with shape (batch_size, height_out, width_out, channels).

Return type:

`mean()`

max.experimental.functional.mean(x, axis=-1)

source

Computes the mean of elements along a specified axis.

Parameters:

x (Tensor) – The input tensor.
axis (int | None) – The axis along which to reduce. When None, the tensor is flattened to 1-D and reduced. Defaults to -1.

Returns:

A Tensor containing the mean along axis. For an integer axis, it has the same rank as x 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 for x.

Return type:

`min()`

max.experimental.functional.min(x, y=None, /, axis=-1)

source

Computes the minimum of a tensor, or the element-wise minimum of two tensors.

Called with one argument, reduces x along axis. Called with two tensor arguments, returns their element-wise minimum.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]])

row_min = F.min(x, axis=-1)
# row_min has shape (2, 1): [[0.8], [1.9]]

col_min = F.min(x, axis=0)
# col_min has shape (1, 4): [[1.2, 1.9, 2.1, 0.8]]

y = Tensor([[2.0, 2.0, 2.0, 2.0], [2.0, 2.0, 2.0, 2.0]])
element_wise = F.min(x, y)
# element_wise: [[1.2, 2.0, 2.0, 0.8], [2.0, 1.9, 2.0, 2.0]]

Parameters:

x (Tensor) – The input tensor.
y (Tensor | None) – Optional second tensor. When provided, the result is the element-wise minimum of x and y.
axis (int | None) – The axis to reduce along when y is omitted. When None, the tensor is flattened to 1-D first. Defaults to -1.

Returns:

A Tensor containing either the reduced minimum along axis or the element-wise minimum with the broadcast shape of the inputs.

Raises:

ValueError – If axis is out of range for x when reducing.

Return type:

`mod()`

max.experimental.functional.mod(lhs, rhs)

source

Computes the element-wise modulus of two tensors.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([7.0, 10.0, 15.0])
b = Tensor([3.0, 4.0, 6.0])
result = F.mod(a, b)
# result is [1.0, 2.0, 3.0]

Parameters:

lhs (Tensor | int | float) – The dividend tensor or scalar.
rhs (Tensor | int | float) – The divisor tensor or scalar.

Returns:

A Tensor containing lhs % rhs element-wise.

Return type:

`mul()`

max.experimental.functional.mul(lhs, rhs)

source

Multiplies two tensors element-wise.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([1.0, 2.0, 3.0])
b = Tensor([4.0, 5.0, 6.0])
result = F.mul(a, b)
# result is [4.0, 10.0, 18.0]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A Tensor containing the element-wise products.

Return type:

`negate()`

max.experimental.functional.negate(x)

source

Negates a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([-1.0, 0.0, 2.0])
result = F.negate(x)
# result is [1.0, 0.0, -2.0]

Parameters:: x (Tensor) – The input tensor.
Returns:: A Tensor of the same shape and dtype as x containing the negation of each element of x.
Return type:: Tensor

`non_maximum_suppression()`

max.experimental.functional.non_maximum_suppression(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold, out_dim='num_selected')

source

Filters boxes with high intersection-over-union (IoU).

Applies greedy non-maximum suppression independently per (batch, class) pair. For each pair, the algorithm:

Discards boxes whose score is at or below score_threshold.
Sorts the remaining boxes by score in descending order.
Greedily selects boxes, suppressing any later candidate whose IoU with an already-selected box exceeds iou_threshold.
Stops after max_output_boxes_per_class selections per pair.

Boxes use (y1, x1, y2, x2) corner format. Coordinates may be normalized or absolute, since the op handles both. All inputs must be on CPU.

from max.driver import CPU
from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

device = CPU()
# boxes: (batch, num_boxes, 4); scores: (batch, num_classes, num_boxes).
boxes = Tensor.ones((1, 3, 4), dtype=DType.float32, device=device)
scores = Tensor.ones((1, 1, 3), dtype=DType.float32, device=device)
# Each output row is (batch_index, class_index, box_index), with a
# data-dependent number of rows.
result = F.non_maximum_suppression(
    boxes,
    scores,
    max_output_boxes_per_class=Tensor(2, dtype=DType.int64, device=device),
    iou_threshold=Tensor(0.5, dtype=DType.float32, device=device),
    score_threshold=Tensor(0.0, dtype=DType.float32, device=device),
)
# result has shape (num_selected, 3)

Parameters:

boxes (Tensor) – The input boxes tensor of shape (batch_size, num_boxes, 4), with a float dtype.
scores (Tensor) – The per-class scores of shape (batch_size, num_classes, num_boxes), with the same dtype as boxes.
max_output_boxes_per_class (Tensor) – A scalar int64 tensor giving the maximum number of boxes to select per (batch, class) pair.
iou_threshold (Tensor) – A scalar float tensor giving the IoU suppression threshold.
score_threshold (Tensor) – A scalar float tensor giving the minimum score to consider.
out_dim (str) – The name for the dynamic output dimension, which is the number of selected boxes. Defaults to "num_selected".

Returns:

A Tensor containing the selected boxes, with shape (out_dim, 3) and int64 dtype. Each row is (batch_index, class_index, box_index).

Return type:

`nonzero()`

max.experimental.functional.nonzero(x, out_dim)

source

Returns the indices of all nonzero elements of a tensor.

Each row is the multi-index of one nonzero element, and the rows are generated in row-major order.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[0, 1], [2, 0]])
# Nonzero elements at (0, 1) and (1, 0) produce [[0, 1], [1, 0]].
result = F.nonzero(x, out_dim="nonzero")
# result is [[0, 1], [1, 0]]

Parameters:

x (Tensor) – The input tensor.
out_dim (int | str | Dim | integer | TypedAttr) – The new data-dependent dimension for the number of nonzero elements.

Returns:

A Tensor containing the indices of the nonzero elements of x, with shape [out_dim, x.rank] and int64 dtype.

Raises:

ValueError – If x is scalar, or if x is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

`normal()`

max.experimental.functional.normal(shape=(), mean=0.0, std=1.0, *, dtype=None, device=None)

source

Samples values from a Gaussian (normal) distribution with the given mean and standard deviation.

When device is a DeviceMapping, each Sharded axis draws an independent stream while shards on Replicated axes draw identical values.

from max.experimental import functional as F

result = F.gaussian((2, 3), mean=0.0, std=1.0)
# result is a (2, 3) tensor sampled from a standard normal distribution.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the resulting tensor.
mean (float) – The mean of the distribution. Defaults to 0.0.
std (float) – The standard deviation of the distribution. Defaults to 1.0.
dtype (DType | None) – The data type of the tensor.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement.

Returns:

A Tensor of the requested shape, dtype, and placement with values sampled from Normal(mean, std**2).

Return type:

`normal_like()`

max.experimental.functional.normal_like(like, mean=0.0, std=1.0)

source

Samples Gaussian values matching another tensor’s shape and dtype.

Parameters:

like (Tensor | TensorType | DistributedTensorType) – The template tensor whose shape, dtype, and placement are copied.
mean (float) – The mean of the distribution. Defaults to 0.0.
std (float) – The standard deviation of the distribution. Defaults to 1.0.

Returns:

A tensor matching the shape, dtype, and placement of like, with values sampled from Normal(mean, std**2).

Return type:

`not_equal()`

max.experimental.functional.not_equal(lhs, rhs)

source

Tests element-wise inequality between two tensors.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([1.0, 2.0, 3.0])
b = Tensor([1.0, 5.0, 3.0])
result = F.not_equal(a, b)
# result is [False, True, False]

Parameters:

lhs (Tensor | int | float) – The left-hand side tensor or scalar.
rhs (Tensor | int | float) – The right-hand side tensor or scalar.

Returns:

A Tensor with bool dtype containing the element-wise result of lhs != rhs.

Return type:

`ones()`

max.experimental.functional.ones(shape, *, dtype=None, device=None)

source

Creates a tensor filled with ones.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the resulting tensor.
dtype (DType | None) – The data type. Defaults to float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement.

Returns:

A tensor of the requested shape, dtype, and placement with every element set to 1.

Return type:

`ones_like()`

max.experimental.functional.ones_like(like)

source

Creates a tensor filled with ones, matching another tensor’s shape and dtype.

Parameters:: like (Tensor | TensorType | DistributedTensorType) – The template tensor whose shape, dtype, and placement are copied.
Returns:: A tensor matching the shape, dtype, and placement of like, with every element set to 1.
Return type:: Tensor

`outer()`

max.experimental.functional.outer(lhs, rhs)

source

Computes the outer product of two vectors.

from max.experimental import Tensor
from max.experimental import functional as F

lhs = Tensor([1.0, 2.0, 3.0])
rhs = Tensor([4.0, 5.0])
# Outer product, producing [[4, 5], [8, 10], [12, 15]].
result = F.outer(lhs, rhs)
# result has shape (3, 2)

Parameters:

lhs (Tensor) – The left side of the product. Must be rank 1.
rhs (Tensor) – The right side of the product. Must be rank 1.

Returns:

A Tensor containing the outer product of the two input vectors. It has rank 2, with dimension sizes equal to the number of elements of lhs and rhs respectively.

Raises:

ValueError – If lhs or rhs is not rank 1.

Return type:

`pad()`

max.experimental.functional.pad(input, paddings, mode='constant', value=0)

source

Pads a tensor along every dimension.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1, 2], [3, 4]])

# Pad one element before and after each dimension.
result = F.pad(x, [1, 1, 1, 1])
# [[0, 0, 0, 0], [0, 1, 2, 0], [0, 3, 4, 0], [0, 0, 0, 0]]

Parameters:

input (Tensor) – The tensor to pad.
paddings (Iterable[int]) – The amount to pad. For a tensor of rank N, pass 2*N non-negative integers in the order [before_dim0, after_dim0, before_dim1, after_dim1, ...].
mode (Literal['constant', 'reflect', 'edge']) –

How to fill the padded cells. Supported values:
- "constant": fill using value.
- "reflect": reflect the content across each edge, excluding the boundary element (like numpy.pad with mode='reflect').
- "edge": repeat the nearest boundary element (like numpy.pad with mode='edge').
value (Tensor) – The fill value for mode="constant". Defaults to 0.

Returns:

A Tensor containing the padded input, with the same dtype as input.

Raises:

ValueError – If mode is unsupported, or any padding value is negative.
AssertionError – If the number of padding values isn’t twice the input rank.

Return type:

`per_shard_dispatch()`

max.experimental.functional.per_shard_dispatch(graph_op, args, output_mappings)

source

Runs graph_op once per shard and reassembles distributed outputs.

Parameters:

graph_op (Callable[[...], Any]) – The per-rank graph op to run.
args (tuple[Any, ...]) – Already-redistributed args.
output_mappings (tuple[DeviceMapping, ...]) – One DeviceMapping per output.

Return type:

Any

`permute()`

max.experimental.functional.permute(x, dims)

source

Permutes all dimensions of a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

# x has shape (1, 2, 3).
x = Tensor.ones([1, 2, 3])
# Reorder the dimensions to (2, 0, 1), producing shape (3, 1, 2).
result = F.permute(x, [2, 0, 1])

Parameters:

x (Tensor) – The input tensor to permute.
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 x with its dimensions reordered to match dims. It has the same elements and dtype as x, 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:

`pow()`

max.experimental.functional.pow(lhs, rhs)

source

Raises elements of one tensor to the power of another element-wise.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([2.0, 3.0, 4.0])
b = Tensor([3.0, 2.0, 0.5])
result = F.pow(a, b)
# result is [8.0, 9.0, 2.0]

Parameters:

lhs (Tensor | int | float) – The base tensor or scalar.
rhs (Tensor | int | float) – The exponent tensor or scalar.

Returns:

A Tensor with the broadcast shape containing lhs ** rhs element-wise.

Return type:

`print()`

max.experimental.functional.print(value, name)

source

Prints a tensor to the console.

Parameters:

value (Tensor) – The tensor to print.
name (str) – The name of the tensor.

Return type:

None

`prod()`

max.experimental.functional.prod(x, axis=-1)

source

Computes the product of elements along a specified axis.

Parameters:

x (Tensor) – The input tensor.
axis (int | None) – The axis along which to reduce. When None, the tensor is flattened to 1-D and reduced. Defaults to -1.

Returns:

A Tensor containing the product along axis. For an integer axis, it has the same rank as x 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 for x.

Return type:

`qmatmul()`

max.experimental.functional.qmatmul(encoding, config, lhs, *rhs)

source

Performs matrix multiplication between floating point and quantized tensors.

Quantizes the lhs floating point value to match the encoding of the rhs quantized value, performs the matmul, and then dequantizes the result. Compared to a regular matmul op, this one expects the rhs value to be transposed. For example, if the lhs shape is [32, 64] and the quantized rhs shape is also [32, 64], then the output shape is [32, 32]. That is, this function returns the result from:

dequantize(quantize(lhs) @ transpose(rhs))

The last two dimensions in lhs are treated as matrices and multiplied by rhs (which must be a 2-D tensor). Any remaining dimensions in lhs are broadcast dimensions.

Note

This currently supports Q4_0, Q4_K, Q6_K, and supported GPTQ configurations.

Parameters:

encoding (QuantizationEncoding) – The quantization encoding to use.
config (QuantizationConfig | None) – The quantization config. Pass None for Vroom encodings; a supported configuration is required for GPTQ.
lhs (TensorValue) – The non-quantized, left-hand side of the matmul.
rhs (TensorValue) – The transposed and quantized right-hand side tensor(s).

Returns:

A Tensor containing the dequantized, floating point result.

Raises:

ValueError – If encoding is not a supported quantization encoding.
TypeError – If lhs or rhs has an unsupported dtype or rank.
AssertionError – If GPTQ is selected without a configuration.

Return type:

`range()`

max.experimental.functional.range(start, stop, step=1, out_dim=None, *, dtype=None, device=None)

source

Creates a sequence of evenly spaced values from start to stop.

The sequence begins at start and increments by step, stopping before stop (the upper bound is exclusive).

from max.dtype import DType
from max.experimental import functional as F

result = F.range(0, 5, 1, dtype=DType.float32)
# result holds [0.0, 1.0, 2.0, 3.0, 4.0].

Parameters:

start (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The first value in the sequence. Must be a scalar value.
stop (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The exclusive upper bound. The sequence stops before this value. Must be a scalar value.
step (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The spacing between consecutive values. Must be non-zero. Defaults to 1.
out_dim (int | str | Dim | integer | TypedAttr | None) – The expected length of the output. Required when dynamic scalar tensor inputs prevent static length inference. When omitted, it’s computed from scalar literals.
dtype (DType | None) – The element type of the result tensor. Defaults to float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping. Sharded placement is not supported.

Returns:

A Tensor containing the generated sequence.

Raises:

ValueError – If out_dim is omitted for dynamic scalar inputs, if any input isn’t scalar, if any input isn’t on the CPU, or if device requests a sharded placement.
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:

`rebind()`

max.experimental.functional.rebind(x, shape, message='', layout=None)

source

Rebinds the symbolic shape of a tensor.

Asserts at runtime that the tensor’s dimensions match the new shape. Useful for narrowing dynamic dimensions to specific sizes when you have external knowledge of their values.

Parameters:

x (Tensor) – The input tensor.
shape (ShapeLike) – The new symbolic shape.
message (str) – A message included in the runtime assertion if the shapes don’t match. Defaults to "".
layout (FilterLayout | None) – An optional filter layout to attach to the result. Defaults to None.

Returns:

A tensor with the same data and the new symbolic shape.

Return type:

`reduce_scatter()`

max.experimental.functional.reduce_scatter(t, scatter_axis=0, mesh_axis=0, *, even=True)

source

Reduces a tensor across a mesh axis and scatters the result.

Transitions the tensor’s placement on mesh_axis from Partial to Sharded. Each device contributes to the sum and ends up with one shard of the reduced tensor along scatter_axis.

Parameters:

t (Tensor) – The input distributed tensor.
scatter_axis (int) – The tensor axis along which the reduced result is sharded.
mesh_axis (int) – The mesh axis whose placement changes from Partial to Sharded.
even (bool) – Require an even shard split along scatter_axis. Defaults to True.

Returns:

A tensor with the reduced and re-sharded result.

Return type:

`relu()`

max.experimental.functional.relu(x)

source

Applies the ReLU (Rectified Linear Unit) activation element-wise.

ReLU is defined as relu(x) = max(0, x), meaning negative values are set to zero while positive values are unchanged.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]])
result = F.relu(x)
# result is [[0.0, 0.0, 0.0], [1.0, 2.0, 3.0]]

Parameters:: x (Tensor) – The input to the ReLU computation.
Returns:: A Tensor of the same shape and dtype as x containing x with its negative elements replaced by 0.
Return type:: Tensor

`repeat_interleave()`

max.experimental.functional.repeat_interleave(x, repeats, axis=None, out_dim=None)

source

Repeats each element of a tensor along an axis.

Unlike tile(), which repeats whole blocks, this repeats each element repeats times consecutively.

This op runs on CPU only; a GPU input raises an error.

Note

The functional API currently supports only integer repeats. Use max.graph.ops.repeat_interleave() for per-element tensor repeats.

The examples below use an input containing [[1.0, 2.0], [3.0, 4.0]]:

from max.experimental import Tensor
from max.experimental import functional as F
from max.driver import CPU

input = Tensor([[1.0, 2.0], [3.0, 4.0]], device=CPU())

# Repeat each row twice.
output = F.repeat_interleave(input, repeats=2, axis=0)
# [[1, 2], [1, 2], [3, 4], [3, 4]], shape (4, 2)

# Repeat each column twice.
output = F.repeat_interleave(input, repeats=2, axis=1)
# [[1, 1, 2, 2], [3, 3, 4, 4]], shape (2, 4)

# With no axis, flatten the input first, then repeat each element.
output = F.repeat_interleave(input, repeats=2)
# [1, 1, 2, 2, 3, 3, 4, 4], shape (8,)

Parameters:

x (Tensor) – The input tensor.
repeats (int | TensorValue) – The integer number of times to repeat each element.
axis (int | None) – The axis to repeat along. If None (the default), the input is flattened first.
out_dim (int | str | Dim | integer | TypedAttr | None) – The output size along axis. This is inferred when repeats is an integer.

Returns:

A Tensor containing the input with its elements interleaved.

Raises:

ValueError – If repeats is non-positive, if axis is out of range, or if the input is on a GPU device.

Return type:

`reshape()`

max.experimental.functional.reshape(x, shape)

source

Reshapes a 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.experimental import Tensor
from max.experimental import functional as F

# x has shape (2, 3).
x = Tensor.ones([2, 3])
# Reshape the same 6 elements into shape (3, 2).
result = F.reshape(x, [3, 2])

Parameters:

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

Returns:

A Tensor containing x 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:

`resize()`

max.experimental.functional.resize(input, shape, interpolation=InterpolationMode.BILINEAR)

source

Resizes a tensor to a given shape using a specified interpolation method.

from max.experimental import Tensor
from max.experimental import functional as F
from max.graph.ops import InterpolationMode

# NCHW input: batch 1, 1 channel, 2x2 spatial.
x = Tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
# Upscale the spatial dimensions to 4x4.
result = F.resize(x, [1, 1, 4, 4], InterpolationMode.BILINEAR)
# result has shape (1, 1, 4, 4)

Parameters:

input (Tensor) – The input tensor to resize. Must be rank 4 in channels-first (NCHW) layout, (batch_size, channels, height, width).
shape (Iterable[int | str | Dim | integer | TypedAttr]) – The desired output shape, of length 4, layout (batch_size, channels, height, width).
interpolation (InterpolationMode) – The interpolation method, given as an InterpolationMode. Defaults to BILINEAR.

Returns:

A Tensor containing the resized tensor with the given shape.

Raises:

ValueError – If input doesn’t have rank 4, or if shape has the wrong number of elements.

Return type:

`resize_bicubic()`

max.experimental.functional.resize_bicubic(input, size)

source

Resizes a tensor using bicubic interpolation.

Produces an output tensor whose dimensions are given by size using a 4x4-pixel Keys/PyTorch (a = -0.75) cubic convolution filter with half-pixel coordinate mapping.

from max.experimental import Tensor
from max.experimental import functional as F

# NCHW input: batch 1, 1 channel, 2x2 spatial.
x = Tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
# Upscale the spatial dimensions to 4x4.
result = F.resize_bicubic(x, [1, 1, 4, 4])
# result has shape (1, 1, 4, 4)

Parameters:

input (Tensor) – The input tensor to resize. Must be rank 4 in channels-first (NCHW) layout, (batch_size, channels, height, width).
size (Iterable[int | str | Dim | integer | TypedAttr]) – The desired output shape, of length 4, (batch_size, channels, height, width).

Returns:

A Tensor containing the resized tensor, with shape size and the same dtype as input.

Raises:

ValueError – If input doesn’t have rank 4, or if size has a different length.

Return type:

`resize_linear()`

max.experimental.functional.resize_linear(input, size, coordinate_transform_mode=0, antialias=False)

source

Resizes a tensor using linear (bilinear) interpolation.

Produces an output tensor whose shape is given by size using separable 1-D linear filters. It resizes any dimension whose size changes, including the batch and channel dimensions. The operation maps output coordinates back to input coordinates according to coordinate_transform_mode.

from max.experimental import Tensor
from max.experimental import functional as F

# NCHW input: batch 1, 1 channel, 2x2 spatial.
x = Tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
# Upscale the spatial dimensions to 4x4.
result = F.resize_linear(x, [1, 1, 4, 4])
# result has shape (1, 1, 4, 4)

Parameters:

input (Tensor) – The input tensor to resize.
size (Iterable[int | str | Dim | integer | TypedAttr]) – The desired output shape. Must have the same rank as input.
coordinate_transform_mode (int) –

How to map an output coordinate to an input coordinate. Allowed values:
- 0 (half_pixel): Default. Shifts by 0.5 before scaling, consistent with most deep learning frameworks.
- 1 (align_corners): Aligns the corner pixels of the input and output so that the first and last coordinates are preserved exactly.
- 2 (asymmetric): Applies no shift, mapping each output coordinate to coordinate / scale.
- 3 (half_pixel_1D): Like half_pixel, except any axis whose output size is 1 maps to coordinate 0.
antialias (bool) – When True, applies an antialiasing filter when downscaling, which reduces aliasing artifacts by widening the tent filter support by 1 / scale. Has no effect when upscaling. Defaults to False.

Returns:

A Tensor containing the resized tensor, with shape size and the same dtype as input.

Raises:

ValueError – If coordinate_transform_mode isn’t 0-3, or if size has a different rank than input.

Return type:

`resize_nearest()`

max.experimental.functional.resize_nearest(input, size, coordinate_transform_mode=0, round_mode=0)

source

Resizes a tensor using nearest-neighbor interpolation.

Produces an output tensor whose dimensions are given by size by selecting the nearest input sample for each output coordinate.

from max.experimental import Tensor
from max.experimental import functional as F

# NCHW input: batch 1, 1 channel, 2x2 spatial.
x = Tensor([[[[1.0, 2.0], [3.0, 4.0]]]])
# Upscale the spatial dimensions to 4x4.
result = F.resize_nearest(x, [1, 1, 4, 4])
# result has shape (1, 1, 4, 4)

Parameters:

input (Tensor) – The input tensor to resize.
size (Iterable[int | str | Dim | integer | TypedAttr]) – The desired output shape. Must have the same rank as input.
coordinate_transform_mode (int) –

How to map an output coordinate to an input coordinate. Allowed values:
- 0 (half_pixel). Default.
- 1 (align_corners).
- 2 (asymmetric).
- 3 (half_pixel_1D).
See resize_linear() for a description of each mode.
round_mode (int) –

How to round the mapped coordinate to select the nearest input sample. Allowed values:
- 0 (HalfDown, the default): ceil(x - 0.5).
- 1 (HalfUp): floor(x + 0.5).
- 2 (Floor): floor(x).
- 3 (Ceil): ceil(x).

Returns:

A Tensor containing the resized tensor, with shape size and the same dtype as input.

Raises:

ValueError – If coordinate_transform_mode isn’t 0-3, round_mode isn’t 0-3, or size has a different rank than input.

Return type:

`rms_norm()`

max.experimental.functional.rms_norm(input, weight, epsilon, weight_offset=0.0, multiply_before_cast=False)

source

Computes root mean square normalization over the last dimension of input.

The output is input / rms(input) * (weight + weight_offset) where rms(x) = sqrt(mean(x ** 2) + epsilon). Reduction runs over the last axis of input and is broadcast back across the leading axes. See Root Mean Square Layer Normalization for the original formulation.

Two variants are supported through weight_offset and multiply_before_cast:

Llama-style (default): weight_offset=0 and multiply_before_cast=False. The normalized input is cast to the output dtype before multiplication by the weight.
Gemma-style: weight_offset=1 and multiply_before_cast=True. The weight is treated as 1 + weight and multiplication runs in the reduction dtype before casting back.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[3.0, 4.0]])
weight = Tensor([1.0, 1.0])
# Llama-style (default).
y_llama = F.rms_norm(x, weight, epsilon=1e-6)
# Gemma-style treats the weight as 1 + weight.
y_gemma = F.rms_norm(
    x, weight, epsilon=1e-6, weight_offset=1.0, multiply_before_cast=True
)

Parameters:

input (Tensor) – The tensor to normalize. Reduction runs over the last axis.
weight (Tensor) – The scale applied after normalization. A 1-D tensor whose shape matches the last dimension of input.
epsilon (float) – A small positive constant added to the mean of squares for numerical stability.
weight_offset (float) – A value added to weight before scaling. Use 1.0 for Gemma-style normalization and 0.0 otherwise. Defaults to 0.0.
multiply_before_cast (bool) – Whether to multiply by the (offset) weight before casting the normalized input back to the output dtype. Llama-style sets this to False. Defaults to False.

Returns:

A Tensor with the same shape and dtype as input.

Raises:

ValueError – If weight does not match the last dimension of input.

Return type:

`roi_align()`

max.experimental.functional.roi_align(input, rois, output_height, output_width, spatial_scale=1.0, sampling_ratio=0.0, aligned=False, mode='AVG')

source

Applies ROI-align pooling.

Extracts fixed-size feature maps from regions of interest (ROIs) using bilinear interpolation.

Parameters:

input (Tensor) – The input tensor in channels-last (NHWC) layout, (batch_size, height, width, channels).
rois (Tensor) – The regions of interest with shape (num_rois, 5), where each row is (batch_index, x1, y1, x2, y2).
output_height (int) – The height of each output feature map.
output_width (int) – The width of each output feature map.
spatial_scale (float) – The multiplicative factor mapping ROI coordinates to input spatial coordinates. Defaults to 1.0.
sampling_ratio (float) – The number of sampling points per bin in each direction. 0 means adaptive (ceil(bin_size)). Defaults to 0.0.
aligned (bool) – When True, applies a half-pixel offset to ROI coordinates for more precise alignment. Defaults to False.
mode (str) – The pooling mode, either "AVG" or "MAX". Defaults to "AVG".

Returns:

A Tensor containing the pooled values, with shape (num_rois, output_height, output_width, channels).

Raises:

ValueError – If input isn’t rank 4, rois isn’t rank 2 with 5 columns, or mode is invalid.

Return type:

`round()`

max.experimental.functional.round(x)

source

Rounds a tensor to the nearest integer element-wise.

Values exactly halfway between two integers round to the nearest even integer (for example, 2.5 rounds to 2.0 and 3.5 rounds to 4.0). All other values follow normal rounding to the nearest integer.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([0.5, 1.5, 2.5, -0.5])
result = F.round(x)
# Ties round to the nearest even integer:
# result is [0.0, 2.0, 2.0, 0.0]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing each element of x rounded to the nearest integer.
Return type:: Tensor

`rsqrt()`

max.experimental.functional.rsqrt(x)

source

Computes the reciprocal square root of a tensor element-wise.

Computes 1 / sqrt(x) for each element.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, 4.0, 16.0])
result = F.rsqrt(x)
# result is [1.0, 0.5, 0.25]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing the reciprocal square root of each element of x.
Return type:: Tensor

`scatter()`

max.experimental.functional.scatter(input, updates, indices, axis=-1)

source

Writes updates into a copy of input at positions given by indices.

from max.experimental import Tensor
from max.experimental import functional as F
from max.driver import CPU
from max.dtype import DType

x = Tensor([1, 2, 3, 4, 5], dtype=DType.int32, device=CPU())
updates = Tensor([10, 20], dtype=DType.int32, device=CPU())
indices = Tensor([0, 3], dtype=DType.int64, device=CPU())
# Overwrite positions 0 and 3, producing [10, 2, 3, 20, 5].
result = F.scatter(x, updates, indices, axis=0)
# result is [10, 2, 3, 20, 5]

Note

When the scatter axis is Sharded, the dispatcher first calls allgather() to make the input Replicated. It doesn’t emit a per-rank-local (Sharded(a_axis), R, R) → Sharded(a_axis) row, because that’s only correct when the caller masks indices and updates per rank. Models that want expert-parallel semantics override scatter.rule with their own rule.

Parameters:

input (Tensor) – The input tensor to write elements to.
updates (Tensor) – A tensor of elements to write to input.
indices (Tensor) – The positions in input to update.
axis (int) – The axis along which indices indexes. Defaults to -1.

Returns:

A Tensor containing input with updates written at indices. It has the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if the input and updates dtypes mismatch, if indices dtype is not int32/int64, if the inputs aren’t all on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
Error – If input, updates, and indices don’t share the same rank, if updates and indices don’t have the same shape, or if any indices dimension exceeds the matching input dimension.

Return type:

`scatter_add()`

max.experimental.functional.scatter_add(input, updates, indices, axis=-1)

source

Creates a new tensor by accumulating updates into input at indices.

Produces an output tensor by scattering elements from updates into input according to indices, summing values at duplicate indices. For a 2-D input with axis=0 the update rule is:

output[indices[i][j]][j] += updates[i][j]

and with axis=1:

output[i][indices[i][j]] += updates[i][j]

Parameters:

input (Tensor) – The input tensor to accumulate into.
updates (Tensor) – A tensor of values to add.
indices (Tensor) – The positions in input to update.
axis (int) – The axis along which indices indexes into. Defaults to -1.

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if the input and updates dtypes mismatch, if indices dtype is not int32/int64, if the inputs aren’t all on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
Error – If input, updates, and indices don’t share the same rank, if updates and indices don’t have the same shape, or if any indices dimension exceeds the matching input dimension.

Return type:

`scatter_max()`

max.experimental.functional.scatter_max(input, updates, indices, axis=-1)

source

Creates a new tensor by scattering the maximum of updates into input.

Produces an output tensor by scattering elements from updates into input according to indices, keeping the maximum at duplicate indices. For a 2-D input with axis=0 the update rule is:

output[indices[i][j]][j] = max(output[indices[i][j]][j], updates[i][j])

and with axis=1:

output[i][indices[i][j]] = max(output[i][indices[i][j]], updates[i][j])

Parameters:

input (Tensor) – The input tensor to scatter into.
updates (Tensor) – A tensor of values to compare.
indices (Tensor) – The positions in input to update.
axis (int) – The axis along which indices indexes into. Defaults to -1.

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if the input and updates dtypes mismatch, if indices dtype is not int32/int64, if the inputs aren’t all on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
Error – If input, updates, and indices don’t share the same rank, if updates and indices don’t have the same shape, or if any indices dimension exceeds the matching input dimension.

Return type:

`scatter_min()`

max.experimental.functional.scatter_min(input, updates, indices, axis=-1)

source

Creates a new tensor by scattering the minimum of updates into input.

Produces an output tensor by scattering elements from updates into input according to indices, keeping the minimum at duplicate indices. For a 2-D input with axis=0 the update rule is:

output[indices[i][j]][j] = min(output[indices[i][j]][j], updates[i][j])

and with axis=1:

output[i][indices[i][j]] = min(output[i][indices[i][j]], updates[i][j])

Parameters:

input (Tensor) – The input tensor to scatter into.
updates (Tensor) – A tensor of values to compare.
indices (Tensor) – The positions in input to update.
axis (int) – The axis along which indices indexes into. Defaults to -1.

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if the input and updates dtypes mismatch, if indices dtype is not int32/int64, if the inputs aren’t all on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
Error – If input, updates, and indices don’t share the same rank, if updates and indices don’t have the same shape, or if any indices dimension exceeds the matching input dimension.

Return type:

`scatter_mul()`

max.experimental.functional.scatter_mul(input, updates, indices, axis=-1)

source

Creates a new tensor by scattering the product of updates into input.

Produces an output tensor by scattering elements from updates into input according to indices, multiplying values at duplicate indices. For a 2-D input with axis=0 the update rule is:

output[indices[i][j]][j] *= updates[i][j]

and with axis=1:

output[i][indices[i][j]] *= updates[i][j]

Parameters:

input (Tensor) – The input tensor to scatter into.
updates (Tensor) – A tensor of values to multiply.
indices (Tensor) – The positions in input to update.
axis (int) – The axis along which indices indexes into. Defaults to -1.

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if the input and updates dtypes mismatch, if indices dtype is not int32/int64, if the inputs aren’t all on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
Error – If input, updates, and indices don’t share the same rank, if updates and indices don’t have the same shape, or if any indices dimension exceeds the matching input dimension.

Return type:

`scatter_nd()`

max.experimental.functional.scatter_nd(input, updates, indices)

source

Scatters slices from updates into a copy of input at N-dimensional indices.

The last dimension of indices is the index vector. Its values select a slice (or scalar) in input. When the index vector length k is less than input.rank, each update writes a whole slice of the trailing input.rank - k dimensions.

from max.experimental import Tensor
from max.experimental import functional as F
from max.driver import CPU
from max.dtype import DType

x = Tensor(
    [[1, 2], [3, 4], [5, 6]], dtype=DType.int32, device=CPU()
)
updates = Tensor(
    [[10, 20], [50, 60]], dtype=DType.int32, device=CPU()
)
indices = Tensor([[0], [2]], dtype=DType.int64, device=CPU())
# Overwrite rows 0 and 2, producing [[10, 20], [3, 4], [50, 60]].
result = F.scatter_nd(x, updates, indices)
# result is [[10, 20], [3, 4], [50, 60]]

Parameters:

input (Tensor) – The input tensor to write elements to.
updates (Tensor) – A tensor of elements to write to input, with shape indices.shape[:-1] + input.shape[k:].
indices (Tensor) – An int32 or int64 tensor specifying where to write updates. Its last dimension k is the index vector length (k <= input.rank) and its leading dimensions may take any shape. Full indexing uses k = input.rank and partial indexing uses k < input.rank.

Returns:

A Tensor containing input with updates scattered in. It has the same shape and dtype as input.

Raises:

ValueError – If dtypes, devices, ranks, or shapes are incompatible, or if indices isn’t an integral tensor.

Return type:

`scatter_nd_add()`

max.experimental.functional.scatter_nd_add(input, updates, indices)

source

Creates a new tensor by accumulating updates into input at N-D indices.

Produces an output tensor by scattering slices from updates into a copy of input according to N-dimensional index vectors, summing values at duplicate index positions. Each index vector is the last dimension of indices and selects a slice (or scalar) in input.

Example for input.shape = [4, 2], indices.shape = [3, 1] (1-D partial indexing, writes whole rows):

output[indices[i, 0], :] += updates[i, :]

Parameters:

input (Tensor) – The input tensor to accumulate into.
updates (Tensor) – A tensor of values to add.
indices (Tensor) – An index tensor whose last dimension is the index vector length k (k <= input.rank).

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If input and updates dtypes mismatch, if indices dtype isn’t int32 or int64, or if the inputs aren’t all on the same device.

Return type:

`scatter_nd_max()`

max.experimental.functional.scatter_nd_max(input, updates, indices)

source

Creates a new tensor by scattering the maximum of updates into input at N-D indices.

Produces an output tensor by scattering slices from updates into a copy of input according to N-dimensional index vectors, keeping the maximum at duplicate index positions. Each index vector is the last dimension of indices and selects a slice (or scalar) in input.

Example for input.shape = [4, 2], indices.shape = [3, 1] (1-D partial indexing, writes whole rows):

output[indices[i, 0], :] = max(output[indices[i, 0], :], updates[i, :])

Parameters:

input (Tensor) – The input tensor to scatter into.
updates (Tensor) – A tensor of values to compare.
indices (Tensor) – An index tensor whose last dimension is the index vector length k (k <= input.rank).

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If input and updates dtypes mismatch, if indices dtype isn’t int32 or int64, or if the inputs aren’t all on the same device.

Return type:

`scatter_nd_min()`

max.experimental.functional.scatter_nd_min(input, updates, indices)

source

Creates a new tensor by scattering the minimum of updates into input at N-D indices.

Produces an output tensor by scattering slices from updates into a copy of input according to N-dimensional index vectors, keeping the minimum at duplicate index positions. Each index vector is the last dimension of indices and selects a slice (or scalar) in input.

Example for input.shape = [4, 2], indices.shape = [3, 1] (1-D partial indexing, writes whole rows):

output[indices[i, 0], :] = min(output[indices[i, 0], :], updates[i, :])

Parameters:

input (Tensor) – The input tensor to scatter into.
updates (Tensor) – A tensor of values to compare.
indices (Tensor) – An index tensor whose last dimension is the index vector length k (k <= input.rank).

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If input and updates dtypes mismatch, if indices dtype isn’t int32 or int64, or if the inputs aren’t all on the same device.

Return type:

`scatter_nd_mul()`

max.experimental.functional.scatter_nd_mul(input, updates, indices)

source

Creates a new tensor by scattering the product of updates into input at N-D indices.

Produces an output tensor by scattering slices from updates into a copy of input according to N-dimensional index vectors, multiplying values at duplicate index positions. Each index vector is the last dimension of indices and selects a slice (or scalar) in input.

Example for input.shape = [4, 2], indices.shape = [3, 1] (1-D partial indexing, writes whole rows):

output[indices[i, 0], :] *= updates[i, :]

Parameters:

input (Tensor) – The input tensor to scatter into.
updates (Tensor) – A tensor of values to multiply.
indices (Tensor) – An index tensor whose last dimension is the index vector length k (k <= input.rank).

Returns:

A Tensor containing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If input and updates dtypes mismatch, if indices dtype isn’t int32 or int64, or if the inputs aren’t all on the same device.

Return type:

`sigmoid()`

max.experimental.functional.sigmoid(x)

source

Applies the sigmoid activation function element-wise.

Computes sigmoid(x) = 1 / (1 + exp(-x)), mapping all values to the range (0, 1).

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]])
result = F.sigmoid(x)
# result is approximately:
# [[0.119, 0.269, 0.5], [0.731, 0.881, 0.953]]

Parameters:: x (Tensor) – The input to the sigmoid computation. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing each element of x mapped to the range (0, 1).
Return type:: Tensor

`silu()`

max.experimental.functional.silu(x)

source

Applies the SiLU (Swish) activation function element-wise.

Computes silu(x) = x * sigmoid(x).

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([-1.0, 0.0, 1.0, 2.0])
result = F.silu(x)
# result is approximately [-0.269, 0.0, 0.731, 1.762]

Parameters:: x (Tensor) – The input to the SiLU computation. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing the SiLU activation applied to each element of x.
Return type:: Tensor

`sin()`

max.experimental.functional.sin(x)

source

Computes the sine of a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([0.0, 0.5, 1.0])
result = F.sin(x)
# result is approximately [0.0, 0.479, 0.841]

Parameters:: x (Tensor) – The input interpreted as radians. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing the sine of each element of x.
Return type:: Tensor

`slice_tensor()`

max.experimental.functional.slice_tensor(x, indices)

source

Slices out a subtensor of the input tensor based on indices.

The semantics of slice_tensor() follow basic NumPy slicing semantics, with one index per dimension. Each index is one of:

An integer.
A scalar tensor (a dynamic integer index).
A slice.
A (slice, out_dim) tuple, which names the output dimension when slicing a dynamic dimension.
None (to insert a size-1 dimension).
Ellipsis (to fill in full slices for the remaining dimensions).

Slice indices must stay within [-dim, dim], and slice steps must be positive.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
# Take rows 0 and 1 and columns 1 and 2, producing [[2, 3], [5, 6]].
result = F.slice_tensor(x, [slice(0, 2), slice(1, 3)])
# result is [[2, 3], [5, 6]]

Parameters:

x (TensorValue) – The input tensor to slice.
indices (SliceIndices) – The per-dimension index expressions. Each entry is an integer, a scalar tensor, a slice, a (slice, out_dim) tuple, None, or Ellipsis.

Returns:

A Tensor containing the sliced subtensor of x.

Raises:

IndexError – If a slice bound or integer index is out of range for its dimension.
ValueError – If x is a scalar, if more indices than dimensions are given, if more than one Ellipsis appears, or if a slice step is 0.
NotImplementedError – If a plain slice targets a dynamic dimension. Pass a (slice, out_dim) tuple instead.

Return type:

`softmax()`

max.experimental.functional.softmax(value, axis=-1)

source

Computes the softmax of a tensor along an axis.

Normalizes the values along axis so that they sum to 1, with each output element representing the exponentiated input divided by the sum of exponentiated values along that axis.

Parameters:

value (Tensor) – The input to the softmax computation. Must have a floating-point dtype.
axis (int) – The axis along which to compute the softmax. Defaults to the final axis (-1).

Returns:

A Tensor of the same shape and dtype as value containing the softmax of value computed along axis.

Return type:

`split()`

max.experimental.functional.split(x, split_size_or_sections, axis=0)

source

Splits a tensor into chunks along an axis.

An int split_size_or_sections produces equal chunks (the last may be smaller); a sequence specifies per-chunk sizes.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
first, second = F.split(x, [2, 4], axis=0)
# first is [1.0, 2.0]
# second is [3.0, 4.0, 5.0, 6.0]

Parameters:

x (Tensor) – The tensor to split.
split_size_or_sections (int | Sequence[int | str | Dim | integer | TypedAttr]) – Either a positive chunk size or a sequence giving the exact size of each output section.
axis (int) – The axis to split. Negative values count from the end. Defaults to 0.

Returns:

A list of tensors in their original order along axis.

Raises:

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

Return type:

`sqrt()`

max.experimental.functional.sqrt(x)

source

Computes the square root of a tensor element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, 4.0, 9.0, 16.0])
result = F.sqrt(x)
# result is [1.0, 2.0, 3.0, 4.0]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype. Negative values produce NaN since MAX doesn’t support complex numbers.
Returns:: A Tensor of the same shape and dtype as x containing the square root of each element of x.
Return type:: Tensor

`squeeze()`

max.experimental.functional.squeeze(x, axis)

source

Removes a dimension of size 1 from a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

# x has shape (2, 1, 3).
x = Tensor.ones([2, 1, 3])
# Remove the size-1 dimension at axis 1, producing shape (2, 3).
result = F.squeeze(x, 1)

Parameters:

x (Tensor) – The input tensor to squeeze.
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 x with the dimension at axis removed. That dimension size must equal 1, so the result holds the same elements as x 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 input.

Return type:

`stack()`

max.experimental.functional.stack(values, axis=0)

source

Stacks tensors along a new axis.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([[1, 2], [3, 4]])
b = Tensor([[5, 6], [7, 8]])

# Stack the two (2, 2) tensors into one (2, 2, 2) tensor
result = F.stack([a, b], axis=0)
# result has shape (2, 2, 2)

Parameters:

values (Iterable[Tensor]) – The tensors to stack. Each must have the same dtype, rank, shape, and device.
axis (int) – The position of the new axis. Negative values count from the end, where -1 inserts the new axis as the last dimension. Defaults to 0.

Returns:

A Tensor containing the stacked inputs. It has one more dimension than the inputs, and the new dimension has size len(values).

Raises:

ValueError – If values is empty, or if the tensors don’t all have the same dtype, rank, shape, and device.
IndexError – If axis is out of range.

Return type:

`sub()`

max.experimental.functional.sub(lhs, rhs)

source

Subtracts two tensors element-wise.

Either operand may be a Python int or float scalar, which is automatically promoted to a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

a = Tensor([10.0, 20.0, 30.0])
b = Tensor([1.0, 2.0, 3.0])
result = F.sub(a, b)
# result is [9.0, 18.0, 27.0]

Parameters:

lhs (Tensor | int | float) – The minuend (left-hand side) tensor or scalar.
rhs (Tensor | int | float) – The subtrahend (right-hand side) tensor or scalar.

Returns:

A Tensor containing the result of lhs - rhs element-wise.

Return type:

`sum()`

max.experimental.functional.sum(x, axis=-1)

source

Computes the sum of elements along a specified axis.

Parameters:

x (Tensor) – The input tensor.
axis (int | None) – The axis along which to reduce. When None, the tensor is flattened to 1-D and reduced. Defaults to -1.

Returns:

A Tensor containing the sum along axis. For an integer axis, it has the same rank as x 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 for x.

Return type:

`tanh()`

max.experimental.functional.tanh(x)

source

Computes the hyperbolic tangent of a tensor element-wise.

This applies tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)), which maps all values to the range (-1, 1).

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]])
result = F.tanh(x)
# result is approximately:
# [[-0.964, -0.762, 0.0], [0.762, 0.964, 0.995]]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing each element of x mapped to the range (-1, 1).
Return type:: Tensor

`tensor_to_layout()`

max.experimental.functional.tensor_to_layout(t)

source

Converts a Tensor to a TensorLayout for sharding-rule evaluation.

t.shape already carries per-device cells on Sharded axes (via PerShardDim), so the rules that fold per-rank cells (notably reshape_rule) can do the correct shape arithmetic directly. Non-distributed tensors fall back to a plain Shape.

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

`tile()`

max.experimental.functional.tile(x, repeats)

source

Repeats a tensor along each of its dimensions.

Each dimension i is copied repeats[i] times, so its output size is x.shape[i] * repeats[i].

This op runs on CPU. An input on another device is copied to CPU for the operation and the result is copied back.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([[1, 2], [3, 4]])

# Repeat the columns twice, leaving the rows unchanged.
result = F.tile(x, [1, 2])
# [[1, 2, 1, 2], [3, 4, 3, 4]]

Parameters:

x (Tensor) – The tensor to tile.
repeats (Iterable[int | str | Dim | integer | TypedAttr]) – The number of copies for each dimension, one positive value per dimension of x.

Returns:

A Tensor containing the tiled input.

Raises:

ValueError – If repeats doesn’t have one value per dimension, if any statically known value isn’t positive, or if x is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

`to_tensors()`

max.experimental.functional.to_tensors(values)

source

Converts graph op results to Tensor, preserving container type.

Recurses one level into list and tuple containers; unknown types pass through unchanged. Returns Tensor for Buffer and TensorValue leaves, and a same-shape container for list/tuple inputs (each leaf converted independently). Any reflects that leaves change type while the container type is preserved.

Parameters:: values (Any)
Return type:: Any

`top_k()`

max.experimental.functional.top_k(input, k, axis=-1)

source

Returns the k largest values along an axis with their indices.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.0, 3.0, 2.0, 5.0, 4.0])
values, indices = F.top_k(x, k=2, axis=-1)
# values is [5, 4] and indices is [3, 4]

Parameters:

input (Tensor) – The input tensor from which to select the top k.
k (int) – The number of values to select from input. Must be in the range [0, input.shape[axis]].
axis (int) – The axis along which to select the top k. Defaults to -1. On a GPU input, only the last axis is supported.

Returns:

A tuple of two Tensor objects. The first holds the top k values along axis, and the second holds their int64 indices in input. Both tensors have the shape of input with the axis dimension reduced to size k.

Return type:

tuple[Tensor, Tensor]

`transfer_to()`

max.experimental.functional.transfer_to(t, target)

source

Moves a tensor to a target device or device mapping.

Handles every kind of placement transition: single-device transfers, scattering an unsharded tensor onto a mesh, redistributing across placements on the same mesh, and gathering then re-distributing across different meshes.

Parameters:

t (Tensor) – The source tensor, distributed or single-device.
target (Device | DeviceMapping | DeviceRef) – A Device to move to a single device, or a DeviceMapping describing the target mesh and placement.

Returns:

A tensor with the requested placement on the target device or mesh.

Return type:

`transpose()`

max.experimental.functional.transpose(x, axis_1, axis_2)

source

Transposes two axes of a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

# x has shape (2, 3).
x = Tensor.ones([2, 3])
# Swap axes 0 and 1, producing shape (3, 2).
result = F.transpose(x, 0, 1)

Parameters:

x (Tensor) – The input tensor to transpose.
axis_1 (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.
axis_2 (int) – The other axis to transpose. If negative, this indexes from the end of the tensor.

Returns:

A Tensor containing the input with axis_1 and axis_2 transposed. It has the same elements and dtype as x, 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 axis_1 or axis_2 is out of range.

Return type:

`trunc()`

max.experimental.functional.trunc(x)

source

Truncates a tensor toward zero element-wise.

from max.experimental import Tensor
from max.experimental import functional as F

x = Tensor([1.5, 2.7, -1.5, -2.7])
result = F.trunc(x)
# result is [1.0, 2.0, -1.0, -2.0]

Parameters:: x (Tensor) – The input tensor. Must have a floating-point dtype.
Returns:: A Tensor of the same shape and dtype as x containing each element of x truncated toward zero.
Return type:: Tensor

`uniform()`

max.experimental.functional.uniform(shape=(), range=(0, 1), *, dtype=None, device=None)

source

Samples values uniformly from the half-open interval [range[0], range[1]).

Values satisfy range[0] <= x < range[1]. When device is a DeviceMapping, each Sharded axis draws an independent stream while shards on Replicated axes draw identical values.

from max.experimental import functional as F

result = F.uniform((2, 3), range=(0.0, 1.0))
# result is a (2, 3) tensor sampled uniformly from [0.0, 1.0).

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the resulting tensor.
range (tuple[float, float]) – A (low, high) pair giving the half-open interval to sample from. Defaults to (0, 1).
dtype (DType | None) – The data type of the tensor.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement.

Returns:

A Tensor of the requested shape, dtype, and placement with values sampled uniformly from [range[0], range[1]).

Return type:

`uniform_like()`

max.experimental.functional.uniform_like(like, range=(0, 1))

source

Samples uniform values matching another tensor’s shape and dtype.

Parameters:

like (Tensor | TensorType | DistributedTensorType) – The template tensor whose shape, dtype, and placement are copied.
range (tuple[float, float]) – A (low, high) pair giving the half-open interval to sample from. Defaults to (0, 1).

Returns:

A tensor matching the shape, dtype, and placement of like, with values sampled uniformly from [range[0], range[1]).

Return type:

`unsqueeze()`

max.experimental.functional.unsqueeze(x, axis)

source

Inserts a dimension of size 1 into a tensor.

from max.experimental import Tensor
from max.experimental import functional as F

# x has shape (3,).
x = Tensor.ones([3])
# Insert a size-1 dimension at axis 0, producing shape (1, 3).
result = F.unsqueeze(x, 0)

Parameters:

x (Tensor) – The input tensor to unsqueeze.
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 x with a new dimension inserted at axis. That dimension has a size of 1, so the result holds the same elements as x with one more dimension.

Raises:

ValueError – If axis is out of bounds.

Return type:

`where()`

max.experimental.functional.where(cond, x, y)

source

Selects elements from two tensors element-wise based on a condition.

At each position, takes the element from x where cond is true and the element from y where it’s false. Scalar x or y operands are promoted to tensors, and the inputs are broadcast to a common shape.

from max.experimental import Tensor
from max.experimental import functional as F
from max.dtype import DType

cond = Tensor([True, False, True], dtype=DType.bool)
x = Tensor([1, 2, 3], dtype=DType.int32)
y = Tensor([10, 20, 30], dtype=DType.int32)
# Take x where True and y where False, producing [1, 20, 3].
result = F.where(cond, x, y)
# result is [1, 20, 3]

Parameters:

cond (Tensor) – The tensor selecting which input to take at each position. Must have a boolean dtype.
x (Tensor | int | float) – The tensor to select from where cond is true.
y (Tensor | int | float) – The tensor to select from where cond is false.

Returns:

A Tensor containing the element-wise selection from x and y according to cond. It has the promoted dtype of x and y, lives on their shared device, and has the broadcast shape of the inputs.

Raises:

ValueError – If cond doesn’t have a boolean dtype, if the inputs aren’t all on the same device, or if the dtypes of x and y can’t be safely promoted.
Error – If the input shapes aren’t broadcast-compatible.

Return type:

`while_loop()`

max.experimental.functional.while_loop(initial_values, predicate, body)

source

Repeatedly executes a body function while a predicate holds.

Both predicate and body receive and return Tensor values. They take the same number and types of arguments as the initial values. The predicate must return a single boolean scalar tensor that controls loop continuation, and that tensor must reside on CPU; the body must return updated values matching the types of initial_values.

from max.driver import CPU
from max.dtype import DType
from max.experimental import Tensor
from max.experimental import functional as F

def predicate(x):
    return x < 10

def body(x):
    return x + 1

x = Tensor(0, dtype=DType.int32, device=CPU())
(result,) = F.while_loop(x, predicate, body)
# Loop continues until ``x >= 10``; result is ``10``.

Parameters:

initial_values (Iterable[Tensor] | Tensor) – The initial values for the loop arguments. Must be non-empty.
predicate (Callable[[...], Tensor]) – A callable that takes the loop arguments and returns a boolean scalar tensor of type bool.
body (Callable[[...], Tensor | Iterable[Tensor]]) – A callable that takes the loop arguments and returns updated values matching the types of initial_values.

Returns:

The output values from the final loop iteration.

Return type:

`zeros()`

max.experimental.functional.zeros(shape, *, dtype=None, device=None)

source

Creates a tensor filled with zeros.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – The shape of the resulting tensor.
dtype (DType | None) – The data type. Defaults to float32 on CPU or bfloat16 on accelerators.
device (Device | DeviceMapping | DeviceRef | None) – A single device or a DeviceMapping for distributed placement.

Returns:

A tensor of the requested shape, dtype, and placement with every element set to 0.

Return type: