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

Python module

max.graph.ops

Implements operations used when staging a graph.

This module provides operations for building a Graph in MAX. Most operations return a TensorValue, which supports standard Python operators such as +, *, and @ (matrix multiplication), as well as convenience methods like reshape() and flatten(). Ops like constant() can also add constant values to your graph.

When an operation receives inputs with different data types (DType), MAX promotes them to a common type before computing the result. To avoid silently widening a type and hurting performance, the common type is always one of the input types; MAX never invents a new, wider type.

To choose between the input types, MAX ranks each one along two axes:

  • Category, ordered bool < unsigned int < signed int < float.
  • Bit width (for example, 8, 16, 32, or 64 bits).

The common type is the input type with the highest category and the largest bit width. For example, promoting int8 and float16 yields float16: float outranks signed int, and 16 bits is wider than 8.

If an input can’t be safely represented in the chosen type, MAX raises an error rather than widening to a different type. For example, promoting uint8 and int8 selects int8 (signed int outranks unsigned int at the same bit width), but int8 can’t represent the largest uint8 values, so MAX raises an error.

When inputs have different shapes, MAX broadcasts them to a common shape. Shapes are aligned from the trailing dimension, and each pair of dimensions must either match exactly, be 1, or be absent. Size-1 dimensions (and any missing leading dimensions) are stretched to match the corresponding dimension of the other input. If the shapes can’t be reconciled under these rules, MAX raises an error.

Operation

class max.graph.ops.Operation

source

Bases: object

asm()

asm(self, *, enable_debug_info: bool = False, pretty_debug_info: bool = False, print_generic_op_form: bool = False, use_local_scope: bool = False, assume_verified: bool = False, skip_regions: bool = False) → str

source

bytecode

property bytecode

source

(self) -> bytes

clone()

clone(self) → max.Operation

source

context

property context

source

(self) -> Context

discardable_attributes

property discardable_attributes

source

(self) -> max.DiscardableAttributes

from_bytecode

from_bytecode = <nanobind.nb_func object>

source

move_after()

move_after(self, arg: max.Operation, /) → None

source

operands

property operands

source

(self) -> Sequence[max.OpOperand]

parent_op

property parent_op

source

(self) -> max.Operation

regions

property regions

source

(self) -> Sequence[mlir::Region]

results

property results

source

(self) -> Sequence[max.Value[max.Type]]

verify()

verify(self, verify_recursively: bool = True) → None

source

abs()

max.graph.ops.abs(x)

source

Computes the absolute value of a tensor element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("abs_example") as graph:
    x = ops.constant([-1.0, 2.0, -3.0], DType.float32, device=device)
    graph.output(ops.abs(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor.

Returns:

A tensor value of the same shape and dtype with each element replaced by its absolute value.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

acos()

max.graph.ops.acos(x)

source

Computes the arccosine of a tensor element-wise.

Returns values in the range [0, π] (radians) for inputs in [-1, 1].

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("acos_example") as graph:
    x = ops.constant(
        [-1.0, 0.0, 0.5, 1.0], DType.float32, device=device
    )
    graph.output(ops.acos(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (TensorValue) – The input tensor with values in [-1, 1]. Values outside this domain are clamped to the valid range. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the arccosine of each element in radians.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

add()

max.graph.ops.add(lhs, rhs)

source

Adds two tensors element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("add_example") as graph:
    lhs = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([4.0, 5.0, 6.0], DType.float32, device=device)
    graph.output(ops.add(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value containing the element-wise sums.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

allgather()

max.graph.ops.allgather(inputs, signal_buffers, axis=0, group_size=None)

source

Collective allgather operation.

This op is a collective op which takes in tensors from different devices and outputs tensors on different devices. In particular, this operation will gather the inputs across different devices and concatenates them along the specified dimension. The result is then broadcasted back to the same devices that the inputs came from.

Parameters:

  • inputs (Iterable[Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray]) – The input tensors to gather.
  • signal_buffers (Iterable[BufferValue | HasBufferValue]) – Device buffer values used for synchronization.
  • axis (int) – Dimension to concatenate the input tensors. Defaults to 0.
  • group_size (int | None) – Optional number of contiguous devices per independent allgather group. Defaults to all devices.

Returns:

An iterable outputs which all hold the gathered output. Each output tensor contains the concatenation of the inputs in its group along the specified dimension.

Return type:

list[TensorValue]

argmax()

max.graph.ops.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.

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

Parameters:

Returns:

A symbolic integer tensor of indices marking the positions of the maximum values along axis. The result has the same rank as x, with the axis dimension reduced to size 1.

Return type:

TensorValue

argmin()

max.graph.ops.argmin(x, axis=-1)

source

Reduces a symbolic tensor using an argmin operation.

When provided with a tensor with all identical elements, on CPU this will return the first element index in the tensor, on GPU this will return an arbitrary index.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor for the operation.
  • axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension. For example, a value of -1 will compute the reduction along the last dimension.

Returns:

A symbolic tensor representing the result of the argmin operation. The tensor will have the same rank as the input tensor, and the same shape except along the axis dimension which will have size 1.

Return type:

TensorValue

argsort()

max.graph.ops.argsort(x, ascending=True)

source

Returns the indices that would sort a tensor.

This function returns the indices that would sort the input tensor along its first dimension. The returned indices are of type int64.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue) – Input tensor to be sorted.
  • ascending (bool) – If True (default), sort in ascending order. If False, sort in descending order.

Returns:

A tensor of indices of the same shape as the input tensor.

Return type:

TensorValue

as_interleaved_complex()

max.graph.ops.as_interleaved_complex(x)

source

Reshapes the input symbolic tensor as complex from alternating (real, imag).

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – A symbolic tensor representing complex numbers as alternating pairs of (real, imag) real-valued numbers. Its last dimension must have an even size.

Returns:

A symbolic tensor representing the complex-valued tensor, but with the values pulled out as complex numbers. The result has the same dimensions for all dimensions except the last dimension, which is halved, and then a final dimension of size 2 representing the complex value.

Return type:

TensorValue

assert_same_device()

max.graph.ops.assert_same_device(*values, **named_values)

source

Raises ValueError if any of the given values are not on the same device.

Parameters:

Return type:

None

atanh()

max.graph.ops.atanh(x)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("atanh_example") as graph:
    x = ops.constant([-0.5, 0.0, 0.5], DType.float32, device=device)
    graph.output(ops.atanh(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor, with values in the range (-1, 1). Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the inverse hyperbolic tangent of each element.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

avg_pool2d()

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

source

Perform a 2D average pooling operation on the input tensor.

Applies a 2D average pooling operation to the input tensor with layout [N, H, W, C]. The pooling operation slides a window of size kernel_size over the spatial dimensions and computes the average value within each window.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor with shape [N, H, W, C].
  • kernel_size (tuple[int | str | Dim | integer | TypedAttr, int | str | Dim | integer | TypedAttr]) – The height and width of the sliding window.
  • stride (int | tuple[int, int]) – The stride of the sliding window. Can be a single integer 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. Can be a single integer 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. Can be a single integer or a tuple (pad_h, pad_w). Defaults to 0.
  • ceil_mode (bool) – If True, uses ceil instead of floor when computing the output spatial shape. Defaults to False.
  • count_boundary (bool) – If True, includes padding elements in the divisor when computing the average. Defaults to True.

Returns:

A symbolic tensor with the average pooling applied, with shape [N, H_out, W_out, C].

Return type:

TensorValue

band_part()

max.graph.ops.band_part(x, num_lower=None, num_upper=None, exclude=False)

source

Masks out everything except a diagonal band of an input matrix.

Copies a tensor setting everything outside the central diagonal band of the matrices to zero, where all but the last two axes are effectively batches, and the last two axes define sub matrices.

Assumes the input has dimensions [I, J, …, M, N], then the output tensor has the same shape as the input, and the values are given by

out[i, j, ..., m, n] = in_band(m, n) * input[i, j,  ..., m, n].

with the indicator function:

in_band(m, n) = ((num_lower is None || (m - n) <= num_lower)) &&
                (num_upper is None || (n - m) <= num_upper))

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor to mask.
  • num_lower (int | None) – The number of diagonal bands to include below the central diagonal. If None, include the entire lower triangle.
  • num_upper (int | None) – The number of diagonal bands to include above the central diagonal. If None, include the entire upper triangle.
  • exclude (bool) – If true, invert the selection of elements to mask. Elements in the band are set to zero.

Returns:

A symbolic tensor value with the configured selection masked out to 0 values, and the remaining values copied from the input tensor.

Raises:

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

Return type:

TensorValue

bottom_k()

max.graph.ops.bottom_k(input, k, axis=-1)

source

Returns tensor with only the bottom K values along given axis.

Parameters:

Returns:

Bottom K values (ascending), Bottom K indices.

Return type:

tuple[TensorValue, TensorValue]

broadcast_to()

max.graph.ops.broadcast_to(x, shape, out_dims=None)

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

import numpy as np
x = ops.constant(np.ones((3, 1)), DType.float32, device=device)
result = ops.broadcast_to(x, [3, 4])
# result has shape (3, 4)

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

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue) – The input symbolic tensor to broadcast. Must not contain any dynamic dimensions.
  • shape (TensorValue | Iterable[int | str | Dim | integer | TypedAttr]) – The target shape. Either a static shape (no dynamic dimensions) or a TensorValue giving the shape at runtime.
  • out_dims (Iterable[int | str | Dim | integer | TypedAttr] | None) – The explicit output dimensions. Required when shape is a TensorValue (used to declare the symbolic output type); ignored otherwise.

Returns:

A symbolic tensor with the same elements as the input but with the target shape.

Raises:

ValueError – If shape is a TensorValue and out_dims is None.

Return type:

TensorValue

buffer_create()

max.graph.ops.buffer_create(type, init_value=None)

source

Creates a new buffer of the given type.

Allocates a fresh BufferValue inside the graph, rather than taking one as a graph input. Use it when a graph needs scratch mutable state that isn’t passed in from outside.

By default the buffer is uninitialized and re-created on every execution. If init_value is provided, the buffer instead becomes persistent state: it is allocated once and filled with init_value a single time when the model is loaded, and the same buffer is reused (and its mutations preserved) across every execution. Use this for a buffer that a kernel mutates in place and only needs zeroed (or otherwise initialized) once, such as a counter that a kernel resets at the end of each call.

The following example creates a buffer and reads it back:

from max.dtype import DType
from max.graph import BufferType, DeviceRef, Graph, ops

with Graph("create_demo") as graph:
    buffer = ops.buffer_create(
        BufferType(DType.float32, shape=[4], device=DeviceRef.CPU())
    )
    graph.output(ops.buffer_load(buffer))

    print(f"shape: {buffer.shape}")  # Output: shape: [Dim(4)]

Parameters:

  • type (BufferType) – The type of the resulting BufferValue.
  • init_value (float | int | bool | None) – An optional scalar to initialize the buffer with once, when the model is loaded. Providing it makes the buffer persistent state that is reused across executions. Must be representable in the buffer’s dtype.

Returns:

A new buffer of the requested type.

Return type:

BufferValue

buffer_load()

max.graph.ops.buffer_load(x)

source

Loads the contents of a buffer into a value-semantic tensor.

Copies the mutable x buffer into a new TensorValue that you can use in value-semantic operations. To write a tensor back into a buffer, use buffer_store().

The following example reads a buffer input into a tensor:

from max.dtype import DType
from max.graph import BufferType, DeviceRef, Graph, ops

buffer_type = BufferType(DType.float32, shape=[2, 2], device=DeviceRef.CPU())

with Graph("load_demo", input_types=[buffer_type]) as graph:
    loaded = ops.buffer_load(graph.inputs[0].buffer)
    graph.output(loaded)

    print(f"shape: {loaded.shape}")  # Output: shape: [Dim(2), Dim(2)]

Parameters:

x (BufferValue) – The buffer to load into a tensor.

Returns:

A new tensor holding a copy of the buffer’s contents.

Return type:

TensorValue

buffer_store()

max.graph.ops.buffer_store(destination, source)

source

Stores a tensor into a buffer, overwriting its contents.

Copies the value-semantic source tensor into the mutable destination buffer in place. Pair it with buffer_load() to read the buffer back. This is how a graph mutates persistent state, such as writing a new entry into a key-value cache.

The following example reads a buffer, adds one, and writes the result back:

from max.dtype import DType
from max.graph import BufferType, DeviceRef, Graph, ops

buffer_type = BufferType(DType.float32, shape=[4], device=DeviceRef.CPU())

with Graph("store_demo", input_types=[buffer_type]) as graph:
    state = graph.inputs[0].buffer
    updated = ops.buffer_load(state) + 1
    ops.buffer_store(state, updated)
    graph.output(updated)

Parameters:

Return type:

None

buffer_store_slice()

max.graph.ops.buffer_store_slice(destination, source, indices)

source

Stores the input tensor to into a slice in the input buffer.

It stores the immutable input tensor source in the mutable tensor destination. This is semantically equivalent to a copy from source tensor to a slice in the destination buffer at index specified by indices.

Parameters:

Return type:

None

call()

max.graph.ops.call(graph, *args, prefix='')

source

Calls a previously defined graph with the provided arguments.

Use this function to invoke a subgraph built with add_subgraph() or build_subgraph(). The primary benefit is that the compiler processes the subgraph definition once, which reduces compile time significantly for models with repeated blocks.

Examples:

Call a subgraph and forward its outputs to the parent graph:

from max.dtype import DType
from max.graph import Graph, ops
from max.graph.type import TensorType, DeviceRef

input_type = TensorType(DType.float32, [10], DeviceRef.CPU())

with Graph("main", input_types=[input_type]) as graph:
    with graph.add_subgraph(
        "add_one", input_types=[input_type]
    ) as sub:
        x = sub.inputs[0].tensor
        one = ops.constant(1, DType.float32, device=DeviceRef.CPU())
        sub.output(ops.elementwise.add(x, one))

    result = ops.call(sub, graph.inputs[0])
    graph.output(*result)

Call a shared subgraph for each layer of a model, resolving different weights at each call site with prefix:

# Build the subgraph once from the first layer.
subgraph = self.layers[0].build_subgraph(
    "transformer_block",
    inputs=h,
    weight_prefix="layers.0.",
)

# Invoke it once per layer with layer-specific weights.
for idx in range(num_layers):
    outputs = ops.call(
        subgraph, *h, prefix=f"layers.{idx}."
    )

Parameters:

  • graph (Graph) – The subgraph to call.
  • *args (Value[Any]) – Arguments to pass to the subgraph. Must match the subgraph’s input types, excluding the chain value (handled internally).
  • prefix (str) – A string prepended to all weight names when the subgraph is invoked. Use this to distinguish repeated calls to the same subgraph. For example, if a transformer block references a weight named attention.wq, calling with prefix="layers.3." resolves it to layers.3.attention.wq in the weights registry. Leave empty if the subgraph contains no placeholder weights.

Returns:

A list of Value objects representing the subgraph’s outputs, excluding any internal chain values.

Return type:

list[Value[Any]]

cast()

max.graph.ops.cast(x, dtype)

source

Casts a symbolic tensor to a different data type.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue) – The input tensor to cast.
  • dtype (DType) – The target dtype to which the tensor is cast.

Returns:

A new symbolic tensor with the same shape as the input and the specified dtype.

Return type:

TensorValue

ceil()

max.graph.ops.ceil(x)

source

Computes the ceil of a tensor element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("ceil_example") as graph:
    x = ops.constant([1.5, -1.5, 2.7, -2.7], DType.float32, device=device)
    graph.output(ops.ceil(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype rounded up toward positive infinity.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

chunk()

max.graph.ops.chunk(x, chunks, axis=0)

source

Chunk the tensor into an exact number of chunks along the specified dim.

Example:

>>> a = TensorValue([1, 2, 3, 4, 5])
>>> chunk(a, 2, 0)
[TensorValue([1, 2]), TensorValue([3, 4])]

Parameters:

Returns:

A list of chunks tensors.

Return type:

list[TensorValue]

concat()

max.graph.ops.concat(original_vals, axis=0)

source

Concatenates a list of symbolic tensors along an axis.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("concat_example") as graph:
    a = ops.constant([[1, 2], [3, 4]], DType.int32, device=device)
    b = ops.constant([[5, 6], [7, 8]], DType.int32, device=device)
    graph.output(
        ops.concat([a, b], axis=0),  # vertical, shape (4, 2)
        ops.concat([a, b], axis=1),  # horizontal, shape (2, 4)
    )

model = InferenceSession().load(graph)
vertical, horizontal = model.execute()

Parameters:

  • original_vals (Iterable[Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray]) – The list of symbolic tensor values to concatenate. All input tensors must have the same rank and the same size in all dimensions except the concatenation axis.
  • axis (int) – The dimension along which to concatenate. Negative values index relative to the end of the tensor shape. For instance, concat(vs, -1) concatenates along the last dimension. Defaults to 0.

Returns:

A symbolic tensor with the same rank, dtype, and device as the inputs, whose size along axis is the sum of the inputs’ sizes along that axis. Every other dimension matches the inputs.

Return type:

TensorValue

cond()

max.graph.ops.cond(pred, out_types, then_fn, else_fn)

source

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

Selects between the then_fn and else_fn branches based on the runtime value of pred. Both branches must return the same number and types of values as specified by out_types. Buffer mutations within a branch are tracked automatically through the chain mechanism.

The predicate is evaluated at runtime to determine which branch to run. Both branches are compiled, but only the selected branch executes.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, TensorType, ops

device = DeviceRef.CPU()

def then_fn():
    return ops.constant(1, DType.int32, device=device)

def else_fn():
    return ops.constant(0, DType.int32, device=device)

with Graph("cond_example") as graph:
    pred = ops.constant(True, DType.bool, device=device)
    graph.output(
        *ops.cond(
            pred,
            [TensorType(DType.int32, [], device=device)],
            then_fn,
            else_fn,
        )
    )

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

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

Raises:

ValueError – If the branches return different numbers of results or if result types don’t match out_types.

Return type:

list[TensorValue]

constant()

max.graph.ops.constant(value, dtype=None, device=None)

source

Adds a node representing a constant operation.

The value of this constant will have the type TensorType with the same shape as value. If value is a scalar type, it will create a TensorType with 0 dimensions.

The constant will be loaded with the specified dtype. If the constant does not fit within the specified dtype, an error is raised.

Warning: Loading the constant could result in precision loss. For example, loading 16777217 as a float32 will result in 16777216.0.

Parameters:

Returns:

A graph value containing the constant data as an attribute.

Return type:

TensorValue

constant_external()

max.graph.ops.constant_external(name, type, align=None)

source

Registers an external constant (weight) in the graph of a given type.

Two external constants with the same name and type refer to the same weight.

Two external constants with the same name and different types are incompatible and will fail compilation.

Parameters:

  • name (str) – The name of the external constant. This should be the fully-qualified weight name and must be unique.
  • type (TensorType) – The type of the constant value.
  • 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 value of the specified type, representing the weight value associated with the name at compile time.

Return type:

TensorValue

conv2d()

max.graph.ops.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.

The op supports 2-D convolution, with the following layout assumptions:

  • input x has NHWC layout, i.e., (batch_size, height, width, in_channels)
  • filter has layout RSCF, i.e., (height, width, in_channels / num_groups, out_channels)
  • 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 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.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("conv2d_example") as graph:
    # NHWC input: batch 1, 2x2 spatial, 1 channel.
    x = ops.constant(
        [[[[1.0], [2.0]], [[3.0], [4.0]]]],
        DType.float32,
        device=device,
    )
    # RSCF filter: 2x2, 1 in-channel, 1 out-channel, all ones.
    filter = ops.constant(
        [[[[1.0]], [[1.0]]], [[[1.0]], [[1.0]]]],
        DType.float32,
        device=device,
    )
    graph.output(ops.conv2d(x, filter))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – An NHWC input tensor to perform the convolution upon.
  • filter (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – 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) – Optional 1-D bias of shape (out_channels,).
  • input_layout (ConvInputLayout) – Layout of the input tensor (default NHWC).
  • filter_layout (FilterLayout) – Layout of the filter tensor (default RSCF).

Returns:

A symbolic tensor value with the convolution applied.

Return type:

TensorValue

conv2d_transpose()

max.graph.ops.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, paddings, and groups.

The op supports the transpose (gradient) of convolution, with the following layout assumptions: (note the out_channel is w.r.t. the original convolution)

  • input x has NHWC layout, i.e., (batch_size, height, width, in_channels)
  • filter has layout RSCF, i.e., (kernel_height, kernel_width, out_channels, in_channels)
  • bias has shape (out_channels,)

The padding values are expected to take the form in the form [[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]].

This op effectively computes the gradient of a convolution with respect to its input (as if the original convolution operation had the same filter and hyperparameters as this op). A visualization of the computation can be found in https://d2l.ai/chapter_computer-vision/transposed-conv.html.

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 input. In 2D ConvTranspose, dim1 here represents H_out and dim2 represents W_out. In python like syntax, padding a 2x4 spatial output with [0, 1, 2, 1] would yield:

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

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

Building a deconvolution graph (filter is RSCF, with out_channels and in_channels w.r.t. the original convolution):

from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("conv2d_transpose_example") as graph:
    # NHWC input: batch 1, 1x1 spatial, 1 channel.
    x = ops.constant([[[[3.0]]]], DType.float32, device=device)
    # RSCF filter: 2x2 kernel, 1 out-channel, 1 in-channel, all ones.
    filter = ops.constant(
        [[[[1.0]], [[1.0]]], [[[1.0]], [[1.0]]]],
        DType.float32,
        device=device,
    )
    graph.output(ops.conv2d_transpose(x, filter))

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – An NHWC input tensor to perform the deconvolution upon.
  • filter (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The convolution filter in RSCF layout: (height, width, out_channels, in_channels).
  • stride (tuple[int, int]) – The stride of the sliding window for each dimension of input. If a single value is given it is replicated in the H and W dimension. By default the N and C dimensions are set to 0.
  • dilation (tuple[int, int]) – The spacing between the kernel points.
  • padding (tuple[int, int, int, int]) – The amount of padding applied to the input.
  • output_paddings (tuple[int, int]) – this argument is meant to resolve the ambiguity of multiple potential output shapes when any stride is greater than 1. Basically, we’ll add output_paddings[i] number of zeros at the end of output’s ith axis. We only support output_paddings = 0.
  • bias (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – Tensor of shape (out_channels,).
  • input_layout (ConvInputLayout) – Layout of the input tensor (default NHWC).
  • filter_layout (FilterLayout) – Layout of the filter tensor (default RSCF).

Returns:

A symbolic tensor value with the convolution applied.

Return type:

TensorValue

conv3d()

max.graph.ops.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, strides, dilations, paddings, and groups.

The op supports 3-D convolution, with the following layout assumptions:

  • input has NDHWC layout, i.e., (batch_size, depth, height, width, in_channels)
  • filter has layout RSCF, i.e., (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 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.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – An NDHWC input tensor to perform the convolution upon.
  • filter (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The convolution filter in RSCF 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) – Optional 1-D bias of shape (out_channels,).
  • input_layout (ConvInputLayout) – Layout of the input tensor (default NDHWC).
  • filter_layout (FilterLayout) – Layout of the filter tensor (default QRSCF).

Returns:

A symbolic tensor value with the convolution applied. Output shape = (batch_size, depth, height, width, out_channels).

Return type:

TensorValue

cos()

max.graph.ops.cos(x)

source

Computes the cosine of a tensor element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("cos_example") as graph:
    x = ops.constant([0.0, 1.5707, 3.1415], DType.float32, device=device)
    graph.output(ops.cos(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input, interpreted as radians. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the cosine of each element.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

cumsum()

max.graph.ops.cumsum(x, axis=-1, exclusive=False, reverse=False)

source

Computes the cumulative sum of the input tensor along the given axis.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor to sum over.
  • axis (int) – The axis along which to compute the sum. If negative, indexes from the last dimension. For example, a value of -1 will compute the sum along the last dimension.
  • exclusive (bool) – If set, start at 0 and exclude the final element. Otherwise, start with the first element. Said another way, cumsum computes [sum(x[…, :i, …]) for i in range(x.shape[axis])]. If exclusive is set, the bounds are instead range(1, x.shape[axis]).
  • reverse (bool) – If set, start from the end. In other words, the first element will be the total sum, with each element following counting downwards; or [sum(x[…, i:, …]) for i in range(x.shape[axis])].

Returns:

A symbolic tensor representing the result of the cumsum operation. The tensor will have the same type as the input tensor. The computed values will be the cumulative sum of the values along the given axis, according to the specified parameters:

  • if exclusive is set, the first value will be 0, and the last value will be excluded from the sum
  • if reverse is set, the sum will be computed starting at the back of the axis back to the front, rather than front-to-back

Raises:

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

Return type:

TensorValue

custom()

max.graph.ops.custom(name, device, values, out_types, parameters=None)

source

Creates a node to execute a custom graph operation in the graph.

The custom op should be registered by annotating a function with the @extensibility.register decorator.

Parameters:

  • name (str) – The op name provided to @extensibility.register.
  • values (Sequence[Value[Any]]) – The op function’s arguments.
  • out_types (Sequence[Type[Any]]) – The list of op function’s return type.
  • parameters (Mapping[str, bool | int | str | DType] | None) – Dictionary of extra parameters expected by the kernel.
  • device (Device | DeviceRef) – Device that the op is assigned to. This becomes a target parameter to the kernel.

Returns:

Symbolic values representing the outputs of the op in the graph. These correspond 1:1 with the types passed as out_types.

Return type:

list[Value[Any]]

dequantize()

max.graph.ops.dequantize(encoding, quantized)

source

Dequantizes a quantized tensor to floating point.

NOTE: Currently this supports Q4_0, Q4_K, and Q6_K encodings only.

Parameters:

Returns:

The dequantized result (a floating point tensor).

Return type:

TensorValue

distributed_broadcast()

max.graph.ops.distributed_broadcast(input, signal_buffers)

source

Broadcast tensor from source GPU to all GPUs.

This op is a collective operation which broadcasts a tensor from the source GPU (where the input tensor resides) to all participating GPUs. Each GPU receives a copy of the input tensor.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Input tensor to broadcast. The device where this tensor resides becomes the root/source of the broadcast.
  • signal_buffers (Iterable[BufferValue | HasBufferValue]) – Device buffer values used for synchronization. The number of signal buffers determines the number of participating GPUs.

Returns:

List of output tensors, one per device. Each output tensor has the same shape and dtype as the input tensor.

Raises:

ValueError – If signal_buffers is empty, if input tensor device is not found in signal buffer devices, or if devices are not unique.

Return type:

list[TensorValue]

distributed_scatter()

max.graph.ops.distributed_scatter(input_chunks, signal_buffers)

source

Scatter different chunks from root GPU to device groups.

Each DP replica group receives a different input chunk. All TP devices within the same replica get the same chunk. Uses a pull-based approach where each GPU reads its chunk from the root GPU via P2P.

Parameters:

  • input_chunks (Iterable[Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray]) – Input tensors to scatter, one per DP replica. All must reside on the same root device. The number of chunks determines dp_size.
  • signal_buffers (Iterable[BufferValue | HasBufferValue]) – Device buffer values used for synchronization. The number of signal buffers determines the number of participating GPUs (ngpus).

Returns:

List of output tensors, one per device. Each output tensor has the same shape and dtype as its replica’s input chunk.

Raises:

ValueError – If fewer than 2 signal buffers, if input chunks are not on the same device, or if devices are not unique.

Return type:

list[TensorValue]

div()

max.graph.ops.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.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("div_example") as graph:
    lhs = ops.constant(
        [6.0, 10.0, 18.0], DType.float32, device=device
    )
    rhs = ops.constant([2.0, 5.0, 6.0], DType.float32, device=device)
    graph.output(ops.div(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value 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.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

equal()

max.graph.ops.equal(lhs, rhs)

source

Tests element-wise equality between two tensors.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("equal_example") as graph:
    lhs = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([1.0, 5.0, 3.0], DType.float32, device=device)
    graph.output(ops.equal(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is True when lhs == rhs.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

erf()

max.graph.ops.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.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("erf_example") as graph:
    x = ops.constant([-1.0, 0.0, 1.0], DType.float32, device=device)
    graph.output(ops.erf(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input to the error function.

Returns:

A tensor value of the same shape and dtype with the error function applied to each element.

Raises:

Error – If the input is not a tensor.

Return type:

TensorValue

exp()

max.graph.ops.exp(x)

source

Computes the exponential of a tensor element-wise.

Use the exp function to build neural networks with attention mechanisms, activation functions, and probability distributions. exp(x) = e^x, where e is Euler’s number.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("exp_example") as graph:
    x = ops.constant([0.0, 1.0, 2.0], DType.float32, device=device)
    graph.output(ops.exp(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input to the exponential function.

Returns:

A tensor value of the same shape and dtype where each element is e raised to the power of the corresponding input element.

Raises:

Error – If the input does not represent a tensor.

Return type:

TensorValue

flatten()

max.graph.ops.flatten(x, start_dim=0, end_dim=-1)

source

Flattens the specified dims of a symbolic tensor.

The number and order of the elements in the tensor is unchanged. All dimensions from start_dim to end_dim (inclusive) are merged into a single output dim.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic 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 symbolic tensor with the same elements as the input, but with dimensions start_dim through end_dim merged into one.

Raises:

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

Return type:

TensorValue

floor()

max.graph.ops.floor(x)

source

Computes the floor of a tensor element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("floor_example") as graph:
    x = ops.constant([1.5, -1.5, 2.7, -2.7], DType.float32, device=device)
    graph.output(ops.floor(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype rounded down toward negative infinity.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

floor_div()

max.graph.ops.floor_div(lhs, rhs)

source

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

The result is rounded toward negative infinity for all operands, matching Python’s //. Integer operands stay in the integer domain: the divide truncates toward zero, then a floor correction is applied for signed integers (a no-op for unsigned or non-negative operands). Floating-point operands compute floor(lhs / rhs).

Unlike div, integer operands are never promoted to float64. This matters on backends without native 64-bit floating-point support (for example, Apple/Metal GPUs), where an f64 intermediate fails to compile. The // operator is intentionally left on its existing floor(div(...)) path here to keep the blast radius minimal; unifying it onto floor_div is a reasonable follow-up.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("floor_div_example") as graph:
    lhs = ops.constant([7, 10, 18], DType.int32, device=device)
    rhs = ops.constant([2, 5, 6], DType.int32, device=device)
    graph.output(ops.floor_div(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

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

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

fold()

max.graph.ops.fold(input, output_size, kernel_size, stride=1, dilation=1, padding=0)

source

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

The input tensor must have 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.

The resulting output tensor will have shape (N, C, output_shape[0], output_shape[1]).

L, the number of blocks, must be equivalent to: prod((output_size[d] + 2 * padding[d] - dilation[d] * (kernel_size[d] - 1) - 1) / stride[d] + 1)

where d is over all spatial dimensions.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The 3D tensor to fold with shape (N, C * kernel sizes, L).
  • output_size (tuple[int | str | Dim | integer | TypedAttr, int | str | Dim | integer | TypedAttr]) – Spatial dimensions of the output tensor. Must be a tuple of two ints.
  • kernel_size (tuple[int | str | Dim | integer | TypedAttr, int | str | Dim | integer | TypedAttr]) – The size of the sliding blocks. Must be a tuple of two ints.
  • stride (int | tuple[int, int]) – The stride of the sliding blocks in the input dimension (can be an int or a tuple of two ints).
  • dilation (int | tuple[int, int]) – The spacing between the kernel elements. (can be an int or a tuple of two ints).
  • padding (int | tuple[int, int]) – 0-paddings to be added on both sides of the inputs. (can be an int or a tuple of two ints).

Returns:

The folded 4D tensor with shape (N, C, output_shape[0], output_shape[1]).

Return type:

TensorValue

gather()

max.graph.ops.gather(input, indices, axis)

source

Selects elements out of an input tensor by index.

Parameters:

Returns:

A new symbolic tensor representing the result of the gather operation.

Return type:

TensorValue

gather_nd()

max.graph.ops.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.

input_shape = ["a", "b", "c", "d", "e"]
indices_shape = ["a", "f", 3]
input_type = TensorType(DType.bfloat16, input_shape)
indices_type = TensorType(DType.int32, indices_shape)
with Graph("gather_nd", input_types=[input_type, indices_type]) as graph:
    input, indices = graph.inputs
    gathered = ops.gather_nd(input, indices, batch_dims=1)
    print(gathered.type)
# Output: TensorType(dtype=DType.bfloat16, shape=["a", "f", "e"])

In this example:

  • batch_dims is 1, so input and indices share the leading “a” dimension.
  • indices has an additional dimension “f” which becomes part of the output.
  • The last dimension of indices (size 3) is the index vector; each value selects into “b”, “c”, and “d” of input.
  • Since batch_dims (1) + index size (3) < input.rank (5), the remaining dimension “e” is sliced into the output.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic tensor to gather from.
  • indices (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – 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 symbolic tensor 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 of the following:

- `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`.
- The leading `batch_dims` of `input` and `indices`
  don’t match.

Return type:

TensorValue

gelu()

max.graph.ops.gelu(x, approximate='none')

source

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

For approximate == "none", MAX computes the exact GELU function.

For approximate == "tanh", MAX uses the approximation:

gelu(x)=0.5x(1.0+tanh(0.7978845608028654(x+0.044715x3)))gelu(x) = 0.5 * x * (1.0 + tanh(0.7978845608028654 * (x + 0.044715 * x**3)))

For approximate == "quick", MAX uses the approximation:

gelu(x)=sigmoid(1.702x)xgelu(x) = sigmoid(1.702 * x) * x

Parameters:

  • x (TensorValue) – The input to the GELU computation.
  • approximate (str) – One of "none", "tanh", or "quick". Defaults to "none".

Returns:

A tensor value of the same shape and dtype with the GELU activation applied element-wise.

Raises:

  • Error – If the input doesn’t represent a tensor.
  • ValueError – If the approximation method is invalid.

greater()

max.graph.ops.greater(lhs, rhs)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("greater_example") as graph:
    lhs = ops.constant([1.0, 5.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([1.0, 2.0, 4.0], DType.float32, device=device)
    graph.output(ops.greater(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is True when lhs > rhs.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

greater_equal()

max.graph.ops.greater_equal(lhs, rhs)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("greater_equal_example") as graph:
    lhs = ops.constant([1.0, 5.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([1.0, 2.0, 4.0], DType.float32, device=device)
    graph.output(ops.greater_equal(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is True when lhs >= rhs.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

group_norm()

max.graph.ops.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.

group_norm executes only on CUDA/HIP GPU targets, so this example builds the graph but does not run it:

from max.dtype import DType
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("group_norm_example") as graph:
    # Shape (batch=1, channels=4, spatial=1, 1); 2 groups of 2 channels.
    x = ops.constant(
        [[[[1.0]], [[3.0]], [[1.0]], [[3.0]]]],
        DType.float32,
        device=device,
    )
    gamma = ops.constant([1.0, 1.0, 1.0, 1.0], DType.float32, device=device)
    beta = ops.constant([0.0, 0.0, 0.0, 0.0], DType.float32, device=device)
    graph.output(
        ops.group_norm(x, gamma, beta, num_groups=2, epsilon=1e-5)
    )

Parameters:

Returns:

A tensor with the same shape and dtype as input.

Raises:

ValueError – If input has fewer than 2 dimensions.

Return type:

TensorValue

hann_window()

max.graph.ops.hann_window(window_length, device, periodic=True, dtype=float32)

source

Calculate a Hann window for a given length.

Hann window function:

H[n]=1/2[1cos(2pin/(N1))]H[n] = 1/2 [1 - cos(2 * pi * n / (N - 1))]

where N is window_length.

Parameters:

  • window_length (int) – The length of the window.
  • device (DeviceRef) – The device to run the operation on.
  • periodic (bool) – bool flag determines whether the returned window trims off the last duplicate value from the symmetric window and is ready to be used as a periodic window with functions like stft(). hann_window(L, periodic=True) == hann_window(L + 1, periodic=False)[:-1])
  • dtype (DType) – The desired data type of the output tensor.

Returns:

A 1-D tensor of size (window_length,) containing the window.

Raises:

  • ValueError – If window_length is negative.
  • TypeError – If window_length is not an integer.

Return type:

TensorValue

inplace_custom()

max.graph.ops.inplace_custom(name, device, values, out_types=None, parameters=None)

source

Creates a node to execute an in-place custom graph operation in the graph.

The custom op should be registered by annotating a function with the @extensibility.register decorator.

Parameters:

  • name (str) – The op name provided to @extensibility.register.
  • device (Device | DeviceRef) – Device that the op is assigned to. This becomes a target parameter to the kernel.
  • values (Sequence[Value[Any]]) – The op function’s arguments.
  • out_types (Sequence[Type[Any]] | None) – Optional sequence of output types for the op.
  • parameters (dict[str, bool | int | str | DType] | None) – Dictionary of extra parameters expected by the kernel.

Return type:

list[Value[Any]]

irfft()

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

source

Compute the inverse real FFT of the input tensor.

Parameters:

  • input_tensor (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue) – The input tensor to compute the inverse real FFT of.
  • n (int | None) – The size of the output tensor. Must be an int, and cannot be a symbolic Buffer. The input tensor will be padded or truncated to n // 2 + 1 along the specified axis.
  • axis (int) – The axis to compute the inverse real FFT of.
  • normalization (Normalization | str) – The normalization to apply to the output tensor. Can be “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. 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. Otherwise, multiple buffers will be allocated.

Returns:

The inverse real FFT of the input tensor. The shape of the output tensor is the same as the shape of the input tensor, except for the axis that the inverse real FFT is computed over, which is replaced by n.

is_inf()

max.graph.ops.is_inf(x)

source

Tests element-wise whether a tensor contains infinite values.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("is_inf_example") as graph:
    x = ops.constant(
        [1.0, float("inf"), 3.0], DType.float32, device=device
    )
    graph.output(ops.is_inf(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor.

Returns:

A tensor value with bool dtype and the same shape, that is True when the input is positive or negative infinity.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

is_nan()

max.graph.ops.is_nan(x)

source

Tests element-wise whether a tensor contains NaN values.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("is_nan_example") as graph:
    x = ops.constant(
        [1.0, float("nan"), 3.0], DType.float32, device=device
    )
    graph.output(ops.is_nan(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor.

Returns:

A tensor value with bool dtype and the same shape, that is True when the input is NaN.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

layer_norm()

max.graph.ops.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.

For example:

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("layer_norm_example") as graph:
    x = ops.constant([[1.0, 3.0]], DType.float32, device=device)
    gamma = ops.constant([1.0, 1.0], DType.float32, device=device)
    beta = ops.constant([0.0, 0.0], DType.float32, device=device)
    graph.output(ops.layer_norm(x, gamma, beta, epsilon=1e-5))

model = InferenceSession().load(graph)
result = model.execute()[0]
# Each row is normalized to zero mean and unit variance.

Parameters:

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:

TensorValue

log()

max.graph.ops.log(x)

source

Computes the natural logarithm of a tensor element-wise.

The natural logarithm is used in loss functions, normalization, and probability calculations in machine learning. It is the inverse of the exponential function: log(x) returns the value y such that x = e^y, where e is Euler’s number.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("log_example") as graph:
    x = ops.constant(
        [1.0, 2.718, 7.389, 20.0], DType.float32, device=device
    )
    graph.output(ops.log(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Note that log(x) is undefined for x <= 0 on real numbers and complex numbers are not currently supported.

Parameters:

Returns:

A tensor value of the same shape with the natural logarithm applied element-wise.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

log1p()

max.graph.ops.log1p(x)

source

Computes log(1 + x) element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("log1p_example") as graph:
    x = ops.constant([0.0, 1.0, 9.0], DType.float32, device=device)
    graph.output(ops.log1p(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

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

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input to the log computation.

Returns:

A tensor value of the same shape and dtype with log(1 + x) applied to each element.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

logical_and()

max.graph.ops.logical_and(lhs, rhs)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("logical_and_example") as graph:
    lhs = ops.constant([True, True, False], DType.bool, device=device)
    rhs = ops.constant([True, False, True], DType.bool, device=device)
    graph.output(ops.logical_and(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is True when both inputs are True.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

logical_not()

max.graph.ops.logical_not(x)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("logical_not_example") as graph:
    x = ops.constant([True, False, True], DType.bool, device=device)
    graph.output(ops.logical_not(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input boolean tensor.

Returns:

A tensor value with bool dtype and the same shape, with each element negated.

Raises:

Error – If the symbol doesn’t represent a tensor.

Return type:

TensorValue

logical_or()

max.graph.ops.logical_or(lhs, rhs)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("logical_or_example") as graph:
    lhs = ops.constant([True, False, False], DType.bool, device=device)
    rhs = ops.constant([False, True, False], DType.bool, device=device)
    graph.output(ops.logical_or(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is True when at least one input is True.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

logical_xor()

max.graph.ops.logical_xor(lhs, rhs)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("logical_xor_example") as graph:
    lhs = ops.constant([True, False, True], DType.bool, device=device)
    rhs = ops.constant([True, True, False], DType.bool, device=device)
    graph.output(ops.logical_xor(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is``True`` when exactly one input is True.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

logsoftmax()

max.graph.ops.logsoftmax(value, axis=-1)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("logsoftmax_example") as graph:
    x = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    graph.output(ops.logsoftmax(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value of the same shape and dtype with the log-softmax applied along axis.

Raises:

Error – If the input is not a tensor.

Return type:

TensorValue

masked_scatter()

max.graph.ops.masked_scatter(input, mask, updates, out_dim)

source

Creates a new symbolic tensor where the updates are written to input where mask is true.

Parameters:

Returns:

A new symbolic tensor representing the result of the masked_scatter operation.

Return type:

TensorValue

matmul()

max.graph.ops.matmul(lhs, rhs)

source

Computes the matrix product of two tensors.

Use matrix multiplication to implement key building blocks like linear transformations, attention mechanisms, and fully connected layers. You can call matmul() directly or use the @ operator, which calls matmul() implicitly.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("matmul_example") as graph:
    lhs = ops.constant(
        [[1.0, 2.0], [3.0, 4.0]], DType.float32, device=device
    )
    rhs = ops.constant(
        [[5.0, 6.0], [7.0, 8.0]], DType.float32, device=device
    )
    graph.output(ops.matmul(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

The innermost two dimensions of each input are treated as a matrix. In the example above lhs has shape (M, K) = (2, 2) and rhs has shape (K, N) = (2, 2), producing an output of shape (M, N) = (2, 2). The K dimensions must match. Any remaining outer (batch) dimensions are broadcast.

If lhs is 1-D it is reshaped to 1xD, and if rhs is 1-D it is reshaped to Dx1. In both cases, the added size-1 dimensions are removed from the output shape.

Parameters:

Returns:

A tensor value representing the matrix product of lhs and rhs. For 2-D inputs, the output shape is (M, N) where lhs is (M, K) and rhs is (K, N). For higher-dimensional inputs, batch dimensions are preserved and the operation is applied to the last two dimensions of each input.

Return type:

TensorValue

max()

max.graph.ops.max(x, y=None, /, axis=None)

source

Overload for ops.elementwise.max and ops.reduction.max.

  • If two tensors are provided, axis is ignored and returns an elementwise maximum.
  • If one tensor is provided, compute ops.reduction.max on the tensor and axis.

Parameters:

Return type:

TensorValue

max_pool2d()

max.graph.ops.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. The input is in (N, H, W, C) (channels-last) layout.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor with shape (N, H, W, C).
  • kernel_size (tuple[int | str | Dim | integer | TypedAttr, int | str | Dim | integer | TypedAttr]) – 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.

Returns:

A symbolic tensor with shape (N, H_out, W_out, C) containing the max-pooled values.

Return type:

TensorValue

mean()

max.graph.ops.mean(x, axis=-1)

source

Reduces a symbolic tensor using a mean operation.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor for the operation.
  • axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension. For example, a value of -1 will compute the reduction along the last dimension.

Returns:

A symbolic tensor representing the result of the mean operation. The tensor will have the same rank as the input tensor, and the same shape except along the axis dimension which will have size 1.

Return type:

TensorValue

min()

max.graph.ops.min(x, y=None, /, axis=None)

source

Overload for ops.elementwise.min and ops.reduction.min.

  • If two tensors are provided, axis is ignored and returns an elementwise minimum.
  • If one tensor is provided, compute ops.reduction.min on the tensor and axis.

Parameters:

Return type:

TensorValue

mod()

max.graph.ops.mod(lhs, rhs)

source

Computes the element-wise modulus of two tensors.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("mod_example") as graph:
    lhs = ops.constant([10.0, 7.0, 5.0], DType.float32, device=device)
    rhs = ops.constant([3.0, 2.0, 4.0], DType.float32, device=device)
    graph.output(ops.mod(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value containing lhs % rhs element-wise.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

mul()

max.graph.ops.mul(lhs, rhs)

source

Multiplies two tensors element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("mul_example") as graph:
    lhs = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([4.0, 5.0, 6.0], DType.float32, device=device)
    graph.output(ops.mul(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value containing the element-wise products.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

negate()

max.graph.ops.negate(x)

source

Negates a tensor element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("negate_example") as graph:
    x = ops.constant([1.0, -2.0, 3.0], DType.float32, device=device)
    graph.output(ops.negate(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor.

Returns:

A tensor value of the same shape and dtype with each element negated.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

non_maximum_suppression()

max.graph.ops.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:

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

Boxes use [y1, x1, y2, x2] corner format. Coordinates may be normalised or absolute; the op handles both.

Parameters:

Returns:

An int64 tensor of shape [out_dim, 3] where each row is [batch_index, class_index, box_index].

Return type:

TensorValue

nonzero()

max.graph.ops.nonzero(x, out_dim)

source

Returns the indices of all nozero elements in a tensor.

Returns a tensor of indices of the nonzero values in the given tensor. The return value is a 2D tensor of shape [out_dim x rank_in], where out_dim is the number of nonzero elements in the input tensor, and rank_in is the rank of the input tensor. Indices are generated in row-major order.

Parameters:

Returns:

A symbolic tensor of indices

Raises:

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

Return type:

TensorValue

not_equal()

max.graph.ops.not_equal(lhs, rhs)

source

Tests element-wise inequality between two tensors.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("not_equal_example") as graph:
    lhs = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([1.0, 5.0, 3.0], DType.float32, device=device)
    graph.output(ops.not_equal(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value with bool dtype that is True when lhs != rhs.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

outer()

max.graph.ops.outer(lhs, rhs)

source

Computes the outer product of two symbolic vectors.

Parameters:

Returns:

A symbolic tensor representing the outer product of the two input vectors. It will have rank 2, with the dimension sizes being the number of elements of lhs and rhs respectively.

Return type:

TensorValue

pad()

max.graph.ops.pad(input, paddings, mode='constant', value=0)

source

Pads a tensor along every dimension.

Adds padding to the input tensor using the specified padding values and mode.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor to pad.

  • paddings (Iterable[int]) – Sequence of padding values. For a tensor with rank N, paddings must contain 2*N non-negative integers in the order [pad_before_dim0, pad_after_dim0, pad_before_dim1, pad_after_dim1, ...].

  • mode (Literal['constant', 'reflect', 'edge']) –

    The padding mode. Supported values:

    • "constant" - fill padded cells with value.
    • "reflect" - reflect values about the content-region edges (excludes the boundary element, equivalent to numpy.pad with mode='reflect').
    • "edge" - repeat the nearest boundary element (equivalent to numpy.pad with mode='edge').
  • value (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The constant fill value (only used when mode='constant'). Defaults to 0.

Returns:

A symbolic tensor with the same dtype as input, padded along each dimension according to paddings.

Raises:

ValueError – If mode is not one of the supported values, or if any padding value is negative.

Return type:

TensorValue

parallel()

max.graph.ops.parallel(inputs, body_fn, *, buffers=None, chain=None, result_types)

source

Execute a function in parallel for each launch via mo.parallel.

Each input bundle holds one TensorValue per launch. All bundles must have the same launch count. The body receives one representative TensorValue per input bundle (typed like the bundle’s first launch) and yields one TensorValue per output bundle; the runtime re-dispatches the body across all launches.

When buffers are provided (e.g. signal buffers for bundled collectives), the body receives an additional BufferValue argument after the input-bundle representatives. Buffers are flat (one per launch) and not bundled.

When chain is provided, the parallel region is sequenced relative to prior ops and the returned out_chain represents completion of all parallel launches. The in-chain enters the body through a trailing _ChainValue argument (after the bundle representatives and the buffer), and the body must thread it and return the resulting out-chain as the second element of a (tensors, out_chain) tuple.

Parameters:

Returns:

[[t0, t1, ...], ...] per output bundle; if chain is provided, returns (results, out_chain).

Return type:

list[list[TensorValue]] | tuple[list[list[TensorValue]], _ChainValue]

permute()

max.graph.ops.permute(x, dims)

source

Permutes all dimensions of a symbolic tensor.

Parameters:

Returns:

A new symbolic tensor with the dimensions permuted to match the passed in order. It has the same elements and dtype, but the order of the elements is different according to the permutation.

Return type:

TensorValue

pow()

max.graph.ops.pow(lhs, rhs)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("pow_example") as graph:
    lhs = ops.constant([2.0, 3.0, 4.0], DType.float32, device=device)
    rhs = ops.constant([3.0, 2.0, 0.5], DType.float32, device=device)
    graph.output(ops.pow(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

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

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

print()

max.graph.ops.print(value, label='debug_tensor')

source

Prints the value of a tensor or a string during graph execution.

This function is used to output the current value of a tensor and is primarily used for debugging purposes within the context of the Max Engine and its graph execution framework. This is particularly useful to verify the intermediate results of your computations are as expected.

By printing the tensor values, you can visualize the data flowing through the graph, which helps in understanding how the operations are transforming the data.

When labeling the function you can assign the output, making it easier to identify which tensor’s value is being printed, especially when there are multiple print statements in a complex graph.

def add_tensors(a: np.ndarray, b: np.ndarray) -> dict[str, Any]:
    input_type = TensorType(dtype=DType.float32, shape=(1,), device=DeviceRef.CPU())
    with Graph(
        "simple_add_graph", input_types=(input_type, input_type)
    ) as graph:
        lhs, rhs = graph.inputs
        out = ops.add(lhs, rhs)
        ops.print(out, label="addition_output")  # Pass the output tensor here

        graph.output(out)
        print("final graph:", graph)

Parameters:

  • value (str | TensorValue) – The value to print. Can be either a string or a TensorValue.
  • label (str) – A label to identify the printed value. Defaults to debug_tensor.

Return type:

None

prod()

max.graph.ops.prod(x, axis=-1)

source

Reduces a symbolic tensor using a product operation.

Computes the product of elements along a specified axis.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor for the operation.
  • axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension. For example, a value of -1 will compute the reduction along the last dimension.

Returns:

A symbolic tensor representing the result of the product operation. The tensor will have the same rank as the input tensor, and the same shape except along the axis dimension which will have size 1.

Return type:

TensorValue

qmatmul()

max.graph.ops.qmatmul(encoding, config, lhs, *rhs)

source

Performs matrix multiplication between floating point and quantized tensors.

This quantizes the lhs floating point value to match the encoding of the rhs quantized value, performs matmul, and then dequantizes the result. Beware that, 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 2D tensor). Any remaining dimensions in lhs are broadcast dimensions.

NOTE: Currently this supports Q4_0, Q4_K, and Q6_K encodings only.

Parameters:

  • encoding (QuantizationEncoding) – The quantization encoding to use.
  • config (QuantizationConfig | None) – Optional quantization config; required for some encodings (for example, GPTQ).
  • lhs (TensorValue) – The non-quantized, left-hand-side of the matmul.
  • rhs (TensorValue) – The transposed and quantized right-hand-side tensor(s).

Returns:

The dequantized result (a floating point tensor).

Return type:

TensorValue

range()

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

source

Creates a sequence of numbers from start to stop (exclusive) with step.

All arguments are mandatory and must have the same element type.

Note the following restrictions on input values:

  1. step must be non-zero.
  2. stop - start must be zero or have the same sign as step.

Parameters:

Returns:

A symbolic tensor value containing the defined range of values.

Return type:

TensorValue

rebind()

max.graph.ops.rebind(x, shape, message='', layout=None)

source

Rebinds a symbolic tensor to a specified set of dimensions.

This does not mutate the symbolic tensor passed in, but instead adds a runtime assert that the input symbolic shape is equivalent to out_dims shape. For example, if the input tensor shape has dynamic/unknown sizes, this will assert a fixed sizes that may be required for a subsequent operation.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic tensor to rebind.
  • shape (Iterable[int | str | Dim | integer | TypedAttr]) – The symbolic shape to assert for x, as a list of Dim values.
  • message (str) – The message printed if the rebind fails at runtime.
  • layout (FilterLayout | None) – A layout of the weights used by some operations like conv.

Returns:

A symbolic tensor with the same elements and shape as the given tensor, but with the symbolic shape asserted to out_dims.

Return type:

TensorValue

relu()

max.graph.ops.relu(x)

source

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

ReLU is defined as relu(x) = max(0, x): negative values are set to zero while positive values are unchanged. It’s one of the most common activation functions in neural networks because of its computational efficiency and its mitigation of the vanishing gradient problem.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("relu_example") as graph:
    x = ops.constant(
        [[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]],
        DType.float32,
        device=device,
    )
    graph.output(ops.relu(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input to the ReLU computation.

Returns:

A tensor value of the same shape and dtype with negative values replaced by 0.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

repeat_interleave()

max.graph.ops.repeat_interleave(x, repeats, axis=None, out_dim=None)

source

Repeats elements of a tensor along the given dimension.

Modeled after torch.repeat_interleave, with the constraint that

For example, given repeats=2 and the following input:

# Input tensor with shape (2, 2)
input = TensorValue(x)  # Contains [[1.0, 2.0], [3.0, 4.0]]

repeat_interleave with axis=0:

# Output tensor with shape (4, 2)
output = repeat_interleave(input, repeats=2, axis=0)
# Contains [[1.0, 2.0], [1.0, 2.0], [3.0, 4.0], [3.0, 4.0]]

repeat_interleave with axis=1:

# Output tensor with shape (2, 4)
output = repeat_interleave(input, repeats=2, axis=1)
# Contains [[1.0, 1.0, 2.0, 2.0], [3.0, 3.0, 4.0, 4.0]]

repeat_interleave with axis=None (the default):

repeat_interleave with repeats=[2, 3] and axis=0:

repeat_value = TensorValue([2, 3])

# Output tensor with shape (5, 2)
output = repeat_interleave(input, repeats=repeat_value, axis=0)
# Contains [[1.0, 2.0], [1.0, 2.0], [3.0, 4.0], [3.0, 4.0], [3.0, 4.0]]
# Output tensor with shape (8,)
output = repeat_interleave(input, repeats=2)  # axis = None
# Contains [1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0]

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor.
  • repeats (int | TensorValue) – The number of repetitions for each element.
  • axis (int | None) – The dimension along which to repeat values. If axis is not specified or None (the default), flatten the input array and repeat the flattened values.
  • out_dim (int | str | Dim | integer | TypedAttr | None) – Optional symbolic dimension for the output size (for graph validation).

Returns:

A symbolic tensor with the elements interleaved.

Raises:

ValueError – If repeats non-positive or if axis is out of range.

Return type:

TensorValue

reshape()

max.graph.ops.reshape(x, shape)

source

Reshapes a symbolic tensor.

The number and order of the elements in the tensor is unchanged. In other words, if you were to iterate over elements in the tensor by major dimension to minor dimension, the iteration order would stay the same.

If a value of -1 is present in the 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 elements of the reshape.

Parameters:

Returns:

A symbolic tensor with the same elements as the original tensor, but in a new shape. Its symbolic shape is the same as shape.

Raises:

ValueError – if input and target shapes’ number of elements mismatch.

Return type:

TensorValue

resize()

max.graph.ops.resize(input, shape, interpolation=InterpolationMode.BILINEAR)

source

Resize the input tensor to the given shape.

This function resizes a tensor using the specified interpolation method. The tensor is expected to have NCHW format (batch, channels, height, width).

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor to resize. Must have rank 4 in NCHW format.
  • shape (Iterable[int | str | Dim | integer | TypedAttr]) – Desired output shape of length 4 corresponding to (N, C, H, W).
  • interpolation (InterpolationMode) – Desired interpolation enum defined by InterpolationMode. Defaults to InterpolationMode.BILINEAR.

Returns:

A resized tensor with the shape specified by the shape argument.

Raises:

ValueError – If the input doesn’t have rank 4, shape has wrong number of elements, or unsupported interpolation mode is specified.

Return type:

TensorValue

resize_bicubic()

max.graph.ops.resize_bicubic(input, size)

source

Resize a tensor using bicubic interpolation.

Produces an output tensor whose dimensions are given by size using a 4x4-pixel Catmull-Rom (a=-0.75) cubic convolution filter with half_pixel coordinate mapping. Input must be rank-4 NCHW.

Parameters:

Returns:

A new symbolic tensor with shape size and the same dtype as input.

Raises:

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

Return type:

TensorValue

resize_linear()

max.graph.ops.resize_linear(input, size, coordinate_transform_mode=0, antialias=False)

source

Resize a tensor using linear (bilinear) interpolation.

Produces an output tensor whose spatial dimensions are given by size using separable 1-D linear filters. The operation maps output coordinates back to input coordinates according to coordinate_transform_mode.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic tensor to resize.

  • size (Iterable[int | str | Dim | integer | TypedAttr]) – 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:

    • 0half_pixel (default): shifts by 0.5 before scaling, consistent with most deep-learning frameworks.
    • 1align_corners: aligns the corner pixels of input and output so that the first and last coordinates are preserved exactly.
    • 2asymmetric: no shift; equivalent to floor-dividing coordinates by the scale factor.
    • 3half_pixel_1D: like half_pixel but only applied to the last spatial dimension.
  • antialias (bool) – When True, applies an antialiasing filter when the output is smaller than the input (i.e. when downscaling), which reduces aliasing artifacts by widening the tent filter support by 1 / scale. Has no effect when upscaling.

Returns:

A new symbolic tensor with shape size and the same dtype as input.

Raises:

ValueError – If coordinate_transform_mode is not 0-3, or if size has a different rank than input.

Return type:

TensorValue

resize_nearest()

max.graph.ops.resize_nearest(input, size, coordinate_transform_mode=0, round_mode=0)

source

Resize 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.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic tensor to resize.

  • size (Iterable[int | str | Dim | integer | TypedAttr]) – 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:

    • 0half_pixel (default).
    • 1align_corners.
    • 2asymmetric.
    • 3half_pixel_1D.
  • round_mode (int) –

    How to round the mapped coordinate to select the nearest input sample. Allowed values:

    • 0HalfDown (default): ceil(x - 0.5).
    • 1HalfUp: floor(x + 0.5).
    • 2Floor: floor(x).
    • 3Ceil: ceil(x).

Returns:

A new symbolic tensor with shape size and the same dtype as input.

Raises:

ValueError – If coordinate_transform_mode is not 0-3, round_mode is not 0-3, or size has a different rank than input.

Return type:

TensorValue

rms_norm()

max.graph.ops.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.

For example:

import numpy as np

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("rms_norm_example") as graph:
    x = ops.constant([[3.0, 4.0]], DType.float32, device=device)
    weight = ops.constant([1.0, 1.0], DType.float32, device=device)
    y_llama = ops.rms_norm(x, weight, epsilon=1e-6)
    y_gemma = ops.rms_norm(
        x, weight, epsilon=1e-6,
        weight_offset=1.0, multiply_before_cast=True,
    )
    graph.output(y_llama, y_gemma)

model = InferenceSession().load(graph)
llama, gemma = model.execute()
assert np.allclose(llama.to_numpy(), [[0.848528, 1.131371]], atol=1e-4)
# weight_offset adds 1.0 to the weight, doubling the result here.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The tensor to normalize. Reduction runs over the last axis.
  • weight (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – 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:

TensorValue

roi_align()

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

source

Perform ROI Align pooling on the input tensor.

Extracts fixed-size feature maps from regions of interest (ROIs) in the input tensor using bilinear interpolation. The input is expected in NHWC layout.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor with shape [N, H, W, C].
  • rois (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – Regions of interest with shape [M, 5], where each row is [batch_index, x1, y1, x2, y2].
  • output_height (int) – Height of each output feature map.
  • output_width (int) – Width of each output feature map.
  • spatial_scale (float) – Multiplicative factor mapping ROI coordinates to input spatial coordinates. Defaults to 1.0.
  • sampling_ratio (float) – Number of sampling points per bin in each direction. 0 means adaptive (ceil(bin_size)). Defaults to 0.0.
  • aligned (bool) – If True, applies a half-pixel offset to ROI coordinates for more precise alignment. Defaults to False.
  • mode (str) – Pooling mode, either "AVG" or "MAX". Defaults to "AVG".

Returns:

A symbolic tensor with shape [M, output_height, output_width, C].

Raises:

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

Return type:

TensorValue

round()

max.graph.ops.round(x)

source

Rounds a tensor to the nearest integer element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("round_example") as graph:
    x = ops.constant([1.5, 2.5, 3.5, -1.5], DType.float32, device=device)
    graph.output(ops.round(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype rounded to the nearest integer.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

rsqrt()

max.graph.ops.rsqrt(x)

source

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("rsqrt_example") as graph:
    x = ops.constant([1.0, 4.0, 9.0, 16.0], DType.float32, device=device)
    graph.output(ops.rsqrt(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the reciprocal square root of each element.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

scatter()

max.graph.ops.scatter(input, updates, indices, axis=-1)

source

Creates a new symbolic tensor where the updates are written to input according to indices.

Parameters:

Returns:

A new symbolic tensor representing the result of the scatter operation.

Raises:

ValueError – If axis is out of range, if dtypes mismatch, if indices dtype is not int32/int64, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

TensorValue

scatter_add()

max.graph.ops.scatter_add(input, updates, indices, axis=-1)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if dtypes mismatch, if indices dtype is not int32/int64, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

TensorValue

scatter_max()

max.graph.ops.scatter_max(input, updates, indices, axis=-1)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if dtypes mismatch, if indices dtype is not int32/int64, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

TensorValue

scatter_min()

max.graph.ops.scatter_min(input, updates, indices, axis=-1)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if dtypes mismatch, if indices dtype is not int32/int64, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

TensorValue

scatter_mul()

max.graph.ops.scatter_mul(input, updates, indices, axis=-1)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Raises:

ValueError – If axis is out of range, if dtypes mismatch, if indices dtype is not int32/int64, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.

Return type:

TensorValue

scatter_nd()

max.graph.ops.scatter_nd(input, updates, indices)

source

Creates a new symbolic tensor where the updates are scattered into input at specified indices.

Parameters:

Returns:

A new symbolic tensor representing the result of the scatter_nd operation.

Return type:

TensorValue

scatter_nd_add()

max.graph.ops.scatter_nd_add(input, updates, indices)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Return type:

TensorValue

scatter_nd_max()

max.graph.ops.scatter_nd_max(input, updates, indices)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Return type:

TensorValue

scatter_nd_min()

max.graph.ops.scatter_nd_min(input, updates, indices)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Return type:

TensorValue

scatter_nd_mul()

max.graph.ops.scatter_nd_mul(input, updates, indices)

source

Creates a new symbolic 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:

Returns:

A new symbolic tensor with the same shape and dtype as input.

Return type:

TensorValue

shape_to_tensor()

max.graph.ops.shape_to_tensor(shape)

source

Converts a shape to a tensor.

This is useful for using a shape attribute in an op that expects a tensor value.

Parameters:

shape (Iterable[int | str | Dim | integer | TypedAttr]) – the shape attribute of a tensor value.

Returns:

The TensorValue containing the same value as shape.

Return type:

TensorValue

Example:

>>> x = ops.constant(np.zeros((1,)), DType.int64, device=DeviceRef.CPU())
>>> result = ops.stack([
...     x,
...     ops.shape_to_tensor(x.shape),
... ])
TensorValue(dtype=int64, shape=[StaticDim(dim=2), StaticDim(dim=1)])

shard_and_stack()

max.graph.ops.shard_and_stack(inputs, devices, axis=0)

source

Shards a list of input tensors along a specified axis, producing multiple outputs.

This operation takes multiple input tensors, splits each along the specified axis into len(devices) chunks, and returns one output tensor per device. Each output contains the chunks at the corresponding index stacked from all inputs along a new dimension 0.

This is useful for distributing model weights across multiple devices in tensor parallel configurations.

For example, with 2 inputs A and B, axis=0, and 2 devices:

  • Input A shape [10, D], Input B shape [10, D]
  • Output 0: stack([A[0:5], B[0:5]]) -> shape [2, 5, D] on devices[0]
  • Output 1: stack([A[5:10], B[5:10]]) -> shape [2, 5, D] on devices[1]

With axis=1 and 2 devices:

  • Input A shape [D, 10], Input B shape [D, 10]
  • Output 0: stack([A[:, 0:5], B[:, 0:5]]) -> shape [2, D, 5] on devices[0]
  • Output 1: stack([A[:, 5:10], B[:, 5:10]]) -> shape [2, D, 5] on devices[1]

Parameters:

  • inputs (Sequence[Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray]) – A list of symbolic tensors to shard. All tensors must have the same shape, dtype, and device.
  • devices (Sequence[Device | DeviceRef]) – Target devices for each output tensor. The number of devices determines the number of splits. Each output tensor will be placed on the corresponding device. This enables direct host-to-device transfer without intermediate CPU storage.
  • axis (int) – The axis along which to split each input tensor. Defaults to 0. Supports negative indexing (for example, -1 for last axis).

Returns:

A list of len(devices) tensors, each with shape [num_inputs, D0, …, Daxis//len(devices), …, Dn-1] where the input shape is [D0, …, Daxis, …, Dn-1]. Output i contains the stacked chunks at position i from all input tensors, placed on devices[i].

Raises:

ValueError – If inputs list is empty, if devices list is empty, if input tensors don’t have matching shapes, if the dimension size at the axis is not evenly divisible by len(devices), or if axis is out of bounds.

Return type:

list[TensorValue]

side_stream()

max.graph.ops.side_stream(inputs, body_fn, *, result_types, stream_id=1)

source

Run a block of ops on a side device stream via mo.sequence.

The body executes on the device stream selected by stream_id (0 is the default stream), overlapping independent work on the main stream. Inputs map 1:1 to the body’s block arguments and the body returns one value per result_types entry. The graph compiler binds the whole body to a side-stream device-context view and inserts the cross-stream synchronization at the boundary, so callers never manage streams or events directly.

Parameters:

Returns:

One TensorValue per result_types entry.

Return type:

list[TensorValue]

sigmoid()

max.graph.ops.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). The sigmoid function is commonly used for binary classification tasks and as an activation function in neural networks, particularly in output layers for probability prediction.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("sigmoid_example") as graph:
    x = ops.constant(
        [[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]],
        DType.float32,
        device=device,
    )
    graph.output(ops.sigmoid(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (TensorValue) – The input to the sigmoid computation.

Returns:

A tensor value of the same shape and dtype with values in the range (0, 1).

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

silu()

max.graph.ops.silu(x)

source

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

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

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("silu_example") as graph:
    x = ops.constant(
        [-2.0, 0.0, 1.0, 3.0], DType.float32, device=device
    )
    graph.output(ops.silu(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (TensorValue) – The input to the SiLU computation.

Returns:

A tensor value of the same shape and dtype with the SiLU activation applied element-wise.

Raises:

Error – If the input doesn’t represent a tensor.

sin()

max.graph.ops.sin(x)

source

Computes the sine of a tensor element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("sin_example") as graph:
    x = ops.constant([0.0, 1.5707, 3.1415], DType.float32, device=device)
    graph.output(ops.sin(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input interpreted as radians. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the sine of each element.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

slice_tensor()

max.graph.ops.slice_tensor(x, indices)

source

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

The semantics of slice_tensor() follow NumPy slicing semantics with the following restrictions:

  • Slice indices must not index out of [-dim - 1, dim - 1] for negative step, or [-dim, dim] for positive step.
# Reverse a tensor.
slice_tensor(x, [slice(None, None, -1)])
# Unsqueeze the second last dimension of a tensor.
slice_tensor(x, [..., None, slice(None)])

Returns:

The sliced subtensor of x.

Parameters:

Return type:

TensorValue

softmax()

max.graph.ops.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.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("softmax_example") as graph:
    x = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    graph.output(ops.softmax(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value of the same shape and dtype with the softmax applied along axis.

Raises:

Error – If the input doesn’t represent a tensor.

Return type:

TensorValue

split()

max.graph.ops.split(x, split_sizes, axis=0)

source

Splits the input tensor into multiple tensors along a given dimension.

Parameters:

Returns:

A list of tensors with the same length as split_sizes, where each tensor has the same shape as the input except along the split dimension axis, where the size is given by the corresponding element in split_sizes.

Return type:

list[TensorValue]

sqrt()

max.graph.ops.sqrt(x)

source

Computes the square root of a tensor element-wise.

Square root is commonly used in normalization operations, distance calculations, and statistical operations like standard deviation.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("sqrt_example") as graph:
    x = ops.constant([1.0, 4.0, 9.0, 16.0], DType.float32, device=device)
    graph.output(ops.sqrt(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

sqrt requires non-negative inputs for real-valued results. For tensors that may contain negative values, take the absolute value first.

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the square root of each element.

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

squeeze()

max.graph.ops.squeeze(x, axis)

source

Removes a size-1 dimension from a symbolic tensor.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic 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, squeeze(v, -1) squeezes the last dimension.

Returns:

A symbolic tensor with the same number of elements as the input tensor, and whose rank is 1 less than the rank of the input tensor.

Return type:

TensorValue

stack()

max.graph.ops.stack(values, axis=0)

source

Stacks a list of tensors along a new axis.

Parameters:

  • values (Iterable[Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray]) – A list of symbolic tensor values. Each tensor must have the same dtype and rank, and must have the same dimension size for each dimension.
  • axis (int) – The axis to concatenate along. If negative, indexes relative to the end of the tensor shape plus 1. For instance, stack(vs, -1) will create and stack along a new axis as the last dimension, aad stack(vs, -2) will create and stack along a new dimension which is inserted immediately before the last dimension.

Returns:

A new symbolic tensor representing the result of the stack. It will have rank n+1 where n is the rank of each input tensor. Its size on each dimension other than axis will be the same as each input tensors’, with the new axis inserted. Along the new dimension it will have size len(values).

Return type:

TensorValue

sub()

max.graph.ops.sub(lhs, rhs)

source

Subtracts two tensors element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("sub_example") as graph:
    lhs = ops.constant([5.0, 7.0, 9.0], DType.float32, device=device)
    rhs = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    graph.output(ops.sub(lhs, rhs))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

Returns:

A tensor value containing the result of lhs - rhs element-wise.

Raises:

  • Error – If the input shapes are not compatible for broadcasting.
  • Error – If one of the inputs has an unsupported dtype.
  • Error – If the two symbols are parts of different graphs.

Return type:

TensorValue

sum()

max.graph.ops.sum(x, axis=-1)

source

Reduces a symbolic tensor using a sum operation.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor for the operation.
  • axis (int) – The axis along which to compute the reduction. If negative, indexes from the last dimension. For example, a value of -1 will compute the reduction along the last dimension.

Returns:

A symbolic tensor representing the result of the sum operation. The tensor will have the same rank as the input tensor, and the same shape except along the axis dimension which will have size 1.

Return type:

TensorValue

tanh()

max.graph.ops.tanh(x)

source

Computes the hyperbolic tangent of a tensor element-wise.

Defined as tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)), mapping all values to the range (-1, 1). Commonly used as an activation function in recurrent neural networks (RNNs) and as a hidden-layer activation in feedforward networks. Unlike sigmoid (which maps to (0, 1)), tanh is zero-centered, which can help with gradient flow during training.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("tanh_example") as graph:
    x = ops.constant(
        [[-2.0, -1.0, 0.0], [1.0, 2.0, 3.0]],
        DType.float32,
        device=device,
    )
    graph.output(ops.tanh(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with values in the range (-1, 1).

Raises:

Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

tile()

max.graph.ops.tile(x, repeats)

source

Returns a new tensor by tiling the input along each dimension.

The input is copied N_i times on the i-th dimension, where N_i = repeats[i]. The i-th dimension of the output shape is the i-th dimension of the input shape multiplied by N_i.

Parameters:

Returns:

A symbolic tensor whose i-th dimension size equals x.shape[i] * repeats[i].

Raises:

ValueError – If the length of repeats does not match the rank of x, or if any repeat value is not positive. Also raised for GPU inputs when strict_device_placement=DevicePlacementPolicy.Error.

Return type:

TensorValue

top_k()

max.graph.ops.top_k(input, k, axis=-1)

source

Returns tensor with only top K values along given axis.

Parameters:

Returns:

Top K values, Top K indices

Return type:

tuple[TensorValue, TensorValue]

transfer_to()

max.graph.ops.transfer_to(x, device)

source

Inserts a device transfer node into the compiled graph.

Moves x to device at execution time. This is a graph-level operation: it operates on symbolic TensorValue objects during graph tracing and is baked into the compiled graph as an mo.transfer MLIR op.

This is distinct from to(), which is a pre-compilation operation that moves stored weight tensors on the Python host before the graph is built. Use transfer_to when you need to route an activation tensor between devices inside forward() (for example, host-to-device input staging, device-to-host output retrieval, or cross-GPU tensor movement for multi-device models).

Host↔device transfers (CPU↔GPU) use the graph’s immutable root chain so they can be hoisted to model initialization by the optimizer. Device-to-device transfers (GPU↔GPU) join both per-device chains to prevent reordering that would deadlock multi-device collectives. If source and destination device are identical, this is a no-op.

Parameters:

Returns:

A new TensorValue on the specified device.

Return type:

TensorValue

transpose()

max.graph.ops.transpose(x, axis_1, axis_2)

source

Transposes two axes of a symbolic tensor.

For more information, see transpose().

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic 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, transpose(v, -1, -2) transposes the last two axes.
  • axis_2 (int) – The other axis to transpose. May also be negative to index from the end of the tensor.

Returns:

A new symbolic tensor with the two specified axes transposed. It has the same elements and dtype, but the order of the elements is different according to the transposition.

Return type:

TensorValue

trunc()

max.graph.ops.trunc(x)

source

Truncates a tensor toward zero element-wise.

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()
with Graph("trunc_example") as graph:
    x = ops.constant([1.5, -1.5, 2.7, -2.7], DType.float32, device=device)
    graph.output(ops.trunc(x))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor. Must have a floating-point dtype.

Returns:

A tensor value of the same shape and dtype with the fractional part discarded.

Raises:

Error – If the input doesn’t represent tensor or has a non-floating-point dtype.

Return type:

TensorValue

unsqueeze()

max.graph.ops.unsqueeze(x, axis)

source

Inserts a size-1 dimension into a symbolic tensor.

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic 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 1 plus the rank of the tensor. For example, unsqueeze(v, -1) adds a new dimension at the end, and unsqueeze(v, -2) inserts the dimension immediately before the last dimension.

Returns:

A symbolic tensor with the same number of elements as the input tensor, whose rank is 1 larger than the rank of the input tensor. The result’s shape at the axis dimension is a static dimension of size 1.

Return type:

TensorValue

where()

max.graph.ops.where(condition, x, y)

source

Returns element-wise condition ? x : y for input tensors condition, x, and y.

Parameters:

Returns:

A new symbolic tensor holding either values from either x or y, based on the elements in condition.

Return type:

TensorValue

while_loop()

max.graph.ops.while_loop(initial_values, predicate, body)

source

Repeatedly executes a body function while a predicate holds.

Both the predicate and body take the same number and types of arguments as the initial values. The predicate must return a single boolean scalar tensor of type bool that controls loop continuation. The body must return updated values matching the types of the initial value(s).

from max.dtype import DType
from max.engine import InferenceSession
from max.graph import DeviceRef, Graph, ops

device = DeviceRef.CPU()

def predicate(x):
    return x < 10

def body(x):
    return x + 1

with Graph("while_loop_example") as graph:
    x = ops.constant(0, DType.int32, device=device)
    graph.output(*ops.while_loop(x, predicate, body))

model = InferenceSession().load(graph)
result = model.execute()[0]

Parameters:

  • initial_values (Iterable[Value[Any]] | Value[Any]) – The initial values for the loop arguments. Must be non-empty.
  • predicate (Callable[[...], TensorValue]) – A callable that takes the loop arguments and returns a boolean scalar tensor of type bool determining loop continuation.
  • body (Callable[[...], Value[Any] | Iterable[Value[Any]]]) – 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.

Raises:

  • ValueError – If initial_values is empty.
  • NotImplementedError – If any initial value is a buffer rather than a tensor. Buffer operations are not currently supported.

Return type:

list[TensorValue]