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 TensorValue of the same shape and dtype as x representing the absolute value of each element of x.

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.

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]. For the float16, bfloat16, and float32 dtypes, values outside this domain are clamped to the valid range. For float64, they yield NaN. Must have a floating-point dtype.

Returns:

A TensorValue of the same shape and dtype as x representing the arccosine of each element of x. Values range from [0, π] (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 TensorValue representing 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]

allgather_rms_norm()​

max.graph.ops.allgather_rms_norm(inputs, signal_buffers, gammas, epsilon, weight_offset=0.0)

source

Fused all-gather + RMSNorm across devices (bf16 in/out).

All-gathers inputs (per-device [shard_i, cols] row shards) along axis 0 so every device holds the full [sum(shard_i), cols] tensor, then RMSNorms every gathered row in the same launch, consuming it from registers (no global-memory round-trip). The norm is multiply_before_cast=True.

Full-world (no device grouping): outputs are replicated on every device. A gathered row is a verbatim copy, so the residual is a drop-in for allgather along axis 0.

The fuse-vs-fallback threshold is calibrated on AMD (gfx950): at/below it the fused kernel is bit-identical to allgather + rms_norm, above it the two-launch fallback wins. On other backends the block.sum geometry may differ, so the paths agree only to RMSNorm ULP tolerance.

Parameters:

Returns:

A tuple (normed, residual) of two lists, each with one tensor per device: normed[i] is the RMSNorm of the full gathered tensor and residual[i] is the raw gathered tensor itself (the residual stream). Both are the full replicated shape on every device.

Return type:

tuple[list[TensorValue], 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.

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

device = DeviceRef.CPU()
with Graph("argmax_example") as graph:
    x = ops.constant(
        [[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]],
        DType.float32,
        device=device,
    )
    graph.output(ops.argmax(x, axis=-1))  # shape (2, 1): [[1], [2]]

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

Parameters:

Returns:

A TensorValue with int64 dtype representing the indices of the maximum values along axis. The result has the same rank as x, with the axis dimension reduced to size 1.

Raises:

ValueError – If axis is out of range for the input’s rank.

Return type:

TensorValue

argmin()​

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

source

Returns the indices of the minimum values along an axis.

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

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

device = DeviceRef.CPU()
with Graph("argmin_example") as graph:
    x = ops.constant(
        [[1.2, 3.5, 2.1, 0.8], [2.3, 1.9, 4.2, 3.1]],
        DType.float32,
        device=device,
    )
    graph.output(ops.argmin(x, axis=-1))  # shape (2, 1): [[3], [1]]

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 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 computes the reduction along the last dimension. Defaults to -1.

Returns:

A TensorValue with int64 dtype representing the indices of the minimum values along axis. The result has the same rank as x, with the axis dimension reduced to size 1.

Raises:

ValueError – If axis is out of range for the input’s rank.

Return type:

TensorValue

argsort()​

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

source

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

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

device = DeviceRef.CPU()
with Graph("argsort") as graph:
    x = ops.constant([3.0, 1.0, 2.0], DType.float32, device=device)
    # Ascending order visits 1, 2, 3, so the indices are [1, 2, 0].
    graph.output(ops.argsort(x, ascending=True))

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

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue) – The input tensor to sort. Must be rank 1.
  • ascending (bool) – Whether to sort in ascending order. If False, sorts in descending order. Defaults to True.

Returns:

A TensorValue representing the sorting indices, with the same shape as x and int64 dtype.

Raises:

ValueError – If x is not rank 1.

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 TensorValue of the same shape and dtype as x representing the inverse hyperbolic tangent of each element of x.

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

Applies 2D average pooling.

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

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor in channels-last (NHWC) layout, (batch_size, height, width, channels).
  • 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.
  • count_boundary (bool) – When True, includes padding elements in the divisor when computing the average. Defaults to True.

Returns:

A TensorValue representing the averaged values, with shape (batch_size, height_out, width_out, channels).

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 the tensor, setting everything outside the central diagonal band of each matrix to zero. All but the last two axes are treated as batches, and the last two axes define the matrices.

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

device = DeviceRef.CPU()
with Graph("band_part") as graph:
    x = ops.constant(
        [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]],
        DType.float32,
        device=device,
    )
    # Keep the main diagonal and one sub-diagonal, producing
    # [[1, 0, 0], [1, 1, 0], [0, 1, 1]].
    graph.output(ops.band_part(x, num_lower=1, num_upper=0))

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

Returns:

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

Raises:

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

Return type:

TensorValue

bottom_k()​

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

source

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

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

device = DeviceRef.CPU()
with Graph("bottom_k") as graph:
    x = ops.constant(
        [1.0, 3.0, 2.0, 5.0, 4.0], DType.float32, device=device
    )
    # The two smallest values are 1 and 2 at indices 0 and 2.
    values, indices = ops.bottom_k(x, k=2, axis=-1)
    graph.output(values, indices)

model = InferenceSession().load(graph)
values, indices = model.execute()

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor from which to select the bottom k.
  • k (int) – The number of values to select from input. Must be in the range [0, input.shape[axis]].
  • axis (int) – The axis along which to select the bottom k. Defaults to -1.

Returns:

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

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

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

device = DeviceRef.CPU()
with Graph("broadcast_to_example") as graph:
    x = ops.constant(
        [[1.0], [1.0], [1.0]], DType.float32, device=device
    )
    result = ops.broadcast_to(x, [3, 4])
    # result has shape (3, 4)

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

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 TensorValue 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 ceiling 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 TensorValue of the same shape and dtype as x representing each element of x 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

Splits a tensor into equal-sized chunks 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("chunk_example") as graph:
    x = ops.constant([1, 2, 3, 4, 5, 6], DType.int32, device=device)

    # Split into three equal chunks along axis 0
    chunks = ops.chunk(x, 3, axis=0)  # [1, 2], [3, 4], and [5, 6]
    graph.output(*chunks)

model = InferenceSession().load(graph)
a, b, c = model.execute()

Parameters:

Returns:

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

Raises:

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

Return type:

list[TensorValue]

concat()​

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

source

Concatenates 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),  # [[1, 2], [3, 4], [5, 6], [7, 8]]
        ops.concat([a, b], axis=1),  # [[1, 2, 5, 6], [3, 4, 7, 8]]
    )

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 symbolic tensors to concatenate. They must have the same rank and size on every dimension except axis.
  • axis (int) – The axis to concatenate along. Negative values count from the end. Defaults to 0.

Returns:

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

Raises:

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

Return type:

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

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

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

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

model = InferenceSession().load(graph)
result = model.execute()[0]
# result holds [[1.0, 2.0], [3.0, 4.0]].

Parameters:

  • value (DLPackArray | Sequence[float | number[Any] | Sequence[Number | NestedArray]] | float | number[Any]) – The value to embed. A Python scalar, a (nested) sequence of numbers, or an array-like object that supports DLPack, such as a NumPy array.
  • dtype (DType | None) – The constant tensor’s element type. Required when value is a Python scalar or sequence. For an array-like value, defaults to the array’s dtype.
  • device (Device | DeviceRef | None) – The device the constant lives on. Required when value is a Python scalar or sequence. For an array-like value, defaults to the array’s device.

Returns:

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

Raises:

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

Return type:

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

This uses the following layout assumptions:

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

The padding values are expected to take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after…) and represent padding 0’s before and after the indicated spatial dimensions in 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) – An optional 1-D bias of shape (out_channels,).
  • input_layout (ConvInputLayout) – The layout of the input tensor. Defaults to NHWC.
  • filter_layout (FilterLayout) – The layout of the filter tensor. Defaults to RSCF.

Returns:

A TensorValue representing the result of the convolution, with shape (batch_size, height_out, width_out, out_channels).

Raises:

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

Return type:

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, and paddings.

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

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

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

The padding values take the form (pad_dim1_before, pad_dim1_after, pad_dim2_before, pad_dim2_after, ...) and crop that many rows or columns from the borders of the indicated spatial dimensions of the output. In 2-D transposed convolution, dim1 represents H_out and dim2 represents W_out. In Python-like syntax, cropping a 2x4 spatial output with [0, 1, 2, 1] trims 0 rows from the top, 1 row from the bottom, 2 columns from the left, and 1 column from the right, leaving:

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

cropped_output = [
  [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]) – A tuple (stride_h, stride_w). Defaults to (1, 1).
  • dilation (tuple[int, int]) – The spacing between the kernel points.
  • padding (tuple[int, int, int, int]) – The number of rows and columns cropped from the borders of the output spatial dimensions.
  • output_paddings (tuple[int, int]) – The number of zeros added at the end of each output spatial axis. This resolves the ambiguity between multiple output shapes when a stride is greater than 1. Only 0 is supported.
  • bias (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray | None) – An optional tensor of shape (out_channels,).
  • input_layout (ConvInputLayout) – The layout of the input tensor. Defaults to NHWC.
  • filter_layout (FilterLayout) – The layout of the filter tensor. Defaults to RSCF.

Returns:

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

Raises:

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

Return type:

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, bias, strides, dilations, paddings, and groups.

This uses the following layout assumptions:

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

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

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

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

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

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

device = DeviceRef.CPU()
input_type = TensorType(DType.float32, [1, 4, 4, 4, 1], device=device)
filter_type = TensorType(DType.float32, [2, 2, 2, 1, 1], device=device)
with Graph(
    "conv3d", input_types=[input_type, filter_type]
) as graph:
    x, filter = graph.inputs
    # Output shape is (1, 3, 3, 3, 1).
    graph.output(ops.conv3d(x.tensor, filter.tensor))

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

Returns:

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

Raises:

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

Return type:

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 TensorValue of the same shape and dtype as x representing the cosine of each element of x.

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.

Parameters:

Returns:

A TensorValue representing the dequantized, floating point result.

Raises:

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

Return type:

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

Scatters different chunks from a root GPU to device groups.

Each data-parallel replica group receives a different input chunk. All tensor-parallel devices within the same replica get the same chunk. Uses a pull-based approach where each GPU reads its chunk from the root GPU over peer-to-peer transfers.

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

# Data-parallel size 2, tensor-parallel size 2, across 4 GPUs
devices = [DeviceRef.GPU(id=i) for i in range(4)]
signals = Signals(devices)

with Graph(
    "distributed_scatter",
    input_types=[
        # One input chunk per data-parallel replica on the root GPU
        TensorType(dtype=DType.uint32, shape=[5], device=devices[0]),
        TensorType(dtype=DType.uint32, shape=[5], device=devices[0]),
        *signals.input_types(),
    ],
) as graph:
    input_chunks = [graph.inputs[0].tensor, graph.inputs[1].tensor]
    signal_buffers = [inp.buffer for inp in graph.inputs[2:]]
    # Returns one output per GPU.
    outputs = ops.distributed_scatter(input_chunks, signal_buffers)
    graph.output(*outputs)

Parameters:

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

Returns:

A 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 any input is invalid. This includes when there are no input chunks, the input chunks aren’t on the same device, the signal buffer devices aren’t unique, or the root device isn’t among the signal buffer devices.

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 TensorValue with the broadcast shape representing 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 TensorValue with bool dtype representing the element-wise result of 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. Must have a floating-point dtype.

Returns:

A TensorValue of the same shape and dtype as x representing the error function applied to each element of x.

Raises:

Error – If the input is not a tensor or has a non-floating-point dtype.

Return type:

TensorValue

exp()​

max.graph.ops.exp(x)

source

Computes the exponential of a tensor element-wise.

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

from max.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. Must have a floating-point dtype.

Returns:

A TensorValue of the same shape and dtype as x representing e raised to the power of each element of x.

Raises:

Error – If the input does not represent a tensor or has a non-floating-point dtype.

Return type:

TensorValue

flatten()​

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

source

Flattens the specified dimensions of a symbolic tensor.

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

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

device = DeviceRef.CPU()
with Graph("flatten") as graph:
    # x has shape (2, 2, 2).
    x = ops.constant(
        [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]],
        DType.float32,
        device=device,
    )
    # Merge dimensions 1 and 2 into one, producing shape (2, 4).
    graph.output(ops.flatten(x, start_dim=1))

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 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 TensorValue representing the start_dim through end_dim of x merged into one dimension.

Raises:

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

Return type:

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 TensorValue of the same shape and dtype as x representing each element of x 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.

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 TensorValue with the broadcast shape representing 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 local blocks into a larger tensor.

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

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

device = DeviceRef.CPU()
# Shape (N, C * kernel_h * kernel_w, L) = (1, 1 * 2 * 2, 9).
input_type = TensorType(DType.float32, [1, 4, 9], device=device)
with Graph("fold", input_types=[input_type]) as graph:
    x = graph.inputs[0].tensor
    # Fold nine 2x2 blocks into a 4x4 image, shape (1, 1, 4, 4).
    graph.output(
        ops.fold(x, output_size=(4, 4), kernel_size=(2, 2))
    )

Parameters:

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

Returns:

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

Raises:

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

Return type:

TensorValue

gather()​

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

source

Selects elements out of an input tensor by index.

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

device = DeviceRef.CPU()
with Graph("gather") as graph:
    x = ops.constant(
        [[1, 2], [3, 4], [5, 6]], DType.int32, device=device
    )
    indices = ops.constant([0, 2], DType.int64, device=device)
    # Select rows 0 and 2, producing [[1, 2], [5, 6]].
    graph.output(ops.gather(x, indices, axis=0))

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

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input symbolic tensor to select elements from.
  • indices (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – A symbolic tensor of int32 or int64 index values on the same device as input.
  • axis (int) – The dimension that indices indexes into input. If negative, indexes relative to the end of the input tensor. For example, gather(input, indices, axis=-1) indexes against the last dimension of input.

Returns:

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

Raises:

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

Return type:

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.

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

device = DeviceRef.CPU()
input_type = TensorType(DType.float32, [2, 3, 4, 5, 6], device)
indices_type = TensorType(DType.int64, [2, 7, 3], device)
with Graph(
    "gather_nd", input_types=[input_type, indices_type]
) as graph:
    input, indices = (value.tensor for value in graph.inputs)
    gathered = ops.gather_nd(input, indices, batch_dims=1)
    graph.output(gathered)
# gathered.shape == [2, 7, 6]

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 TensorValue representing the gathered elements, with the same dtype as input. Its shape is the concatenation of:

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

Raises:

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

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.5 * x * (1.0 + tanh(0.7978845608028654 * (x + 0.044715 * x**3)))

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

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

Parameters:

  • x (TensorValue) – The input to the GELU computation. Must have a floating-point dtype.
  • approximate (str) – One of "none", "tanh", or "quick". Defaults to "none".

Returns:

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

Raises:

  • Error – If the input doesn’t represent a tensor or has a non-floating-point dtype.
  • 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 TensorValue with bool dtype representing the element-wise result of 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 TensorValue with bool dtype representing the element-wise result of 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 TensorValue 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

Computes a Hann window of a given length.

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

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

When N is 0, the result is an empty tensor, and when N is 1, the result is [1]. A periodic window instead computes N + 1 points and drops the last one, which makes it suitable for spectral analysis.

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

device = DeviceRef.CPU()
with Graph("hann_window_example") as graph:
    graph.output(ops.hann_window(4, device, periodic=True))

model = InferenceSession().load(graph)
result = model.execute()[0]
# result holds [0.0, 0.5, 1.0, 0.5].

Parameters:

  • window_length (int) – The number of points in the window.
  • device (DeviceRef) – The device the window lives on.
  • periodic (bool) – Whether to return a periodic window. When True, computes a symmetric window of window_length + 1 points and drops the last point, so that hann_window(L, periodic=True) equals hann_window(L + 1, periodic=False)[:-1]. Defaults to True.
  • dtype (DType) – The output tensor’s data type. Defaults to float32.

Returns:

A 1-D TensorValue of shape (window_length,) representing the window.

Raises:

  • ValueError – If window_length is negative.
  • TypeError – If window_length isn’t 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 TensorValue with bool dtype and the same shape as x, representing an element-wise infinity test. An element is True where x 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 TensorValue with bool dtype and the same shape as x, representing an element-wise NaN test. An element is True where x 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.

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 approximately unit variance.

Parameters:

Returns:

A TensorValue 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.

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

from max.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]

Parameters:

x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input to the log computation. Must contain positive values only. Must have a floating-point dtype.

Returns:

A TensorValue of the same shape and dtype as x representing the natural logarithm of each element of x.

Raises:

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

Return type:

TensorValue

log1p()​

max.graph.ops.log1p(x)

source

Computes log(1 + x) element-wise.

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

from max.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]

Parameters:

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

Returns:

A TensorValue of the same shape and dtype as x representing log(1 + x) for each element of x.

Raises:

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

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 TensorValue with bool dtype representing the element-wise logical AND of lhs and 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

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 TensorValue with bool dtype and the same shape as x, representing the element-wise logical NOT of x.

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 TensorValue with bool dtype representing the element-wise logical OR of lhs and 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

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 TensorValue with bool dtype representing the element-wise logical XOR of lhs and 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

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 TensorValue of the same shape and dtype as value representing the log-softmax of value computed along axis.

Raises:

Error – If the input is not a tensor or has a non-floating-point dtype.

Return type:

TensorValue

masked_scatter()​

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

source

Writes updates into a copy of input at positions where mask is true.

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

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

device = DeviceRef.CPU()
with Graph("masked_scatter") as graph:
    x = ops.constant(
        [[1, 2], [3, 4]], DType.int32, device=device
    )
    mask = ops.constant(
        [[True, False], [False, True]], DType.bool, device=device
    )
    updates = ops.constant([10, 20], DType.int32, device=device)
    # Write into the True positions, producing [[10, 2], [3, 20]].
    graph.output(ops.masked_scatter(x, mask, updates, out_dim="num_updates"))

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

Parameters:

Returns:

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

Raises:

ValueError – If the dtypes of input and updates mismatch, or if input and updates are on different devices.

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.

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor in channels-last (NHWC) layout, (batch_size, height, width, channels).
  • 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]) – Padding added to both sides of each spatial dimension. Out-of-bounds positions are excluded from the maximum (equivalently, they use the dtype’s minimum value or negative infinity), so padding cannot win over negative input values. Either a single int applied to both spatial dimensions, or a tuple (pad_h, pad_w). Defaults to 0.
  • ceil_mode (bool) – When True, uses ceil instead of floor when computing the output spatial shape. Defaults to False.

Returns:

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

Return type:

TensorValue

mean()​

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

source

Computes the mean of elements along a specified axis.

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

device = DeviceRef.CPU()
with Graph("mean_example") as graph:
    x = ops.constant(
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
        DType.float32,
        device=device,
    )
    graph.output(ops.mean(x, axis=-1))  # shape (2, 1): [[2.0], [5.0]]

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 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 computes the reduction along the last dimension. Defaults to -1.

Returns:

A TensorValue representing the mean along axis. It has the same rank as x, with the axis dimension reduced to size 1.

Raises:

ValueError – If axis is out of range for the input’s rank.

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 TensorValue representing 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 TensorValue representing 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 TensorValue of the same shape and dtype as x representing the negation of each element of x.

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 the 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 normalized or absolute, since the op handles both.

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

device = DeviceRef.CPU()
box_type = TensorType(DType.float32, [1, 3, 4], device=device)
score_type = TensorType(DType.float32, [1, 1, 3], device=device)
with Graph(
    "nms", input_types=[box_type, score_type]
) as graph:
    boxes, scores = graph.inputs
    # Each output row is (batch_index, class_index, box_index), with a
    # data-dependent number of rows.
    selected = ops.non_maximum_suppression(
        boxes.tensor,
        scores.tensor,
        max_output_boxes_per_class=ops.constant(
            2, DType.int64, device=device
        ),
        iou_threshold=ops.constant(0.5, DType.float32, device=device),
        score_threshold=ops.constant(
            0.0, DType.float32, device=device
        ),
    )
    graph.output(selected)

Parameters:

Returns:

A TensorValue representing the selected boxes, with shape (out_dim, 3) and int64 dtype. 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 nonzero elements of a tensor.

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

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

device = DeviceRef.CPU()
with Graph("nonzero") as graph:
    x = ops.constant([[0, 1], [2, 0]], DType.int32, device=device)
    # Nonzero elements at (0, 1) and (1, 0) produce [[0, 1], [1, 0]]
    graph.output(ops.nonzero(x, out_dim="nonzero"))

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

Parameters:

Returns:

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

Raises:

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

Return type:

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 TensorValue with bool dtype representing the element-wise result of 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.

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

device = DeviceRef.CPU()
with Graph("outer") as graph:
    lhs = ops.constant([1.0, 2.0, 3.0], DType.float32, device=device)
    rhs = ops.constant([4.0, 5.0], DType.float32, device=device)
    # Outer product, producing [[4, 5], [8, 10], [12, 15]].
    graph.output(ops.outer(lhs, rhs))

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

Parameters:

Returns:

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

Raises:

ValueError – If lhs or rhs is not rank 1.

Return type:

TensorValue

pad()​

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

source

Pads a tensor along every dimension.

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

device = DeviceRef.CPU()
with Graph("pad_example") as graph:
    x = ops.constant([[1, 2], [3, 4]], DType.int32, device=device)

    # Add a dim before and after dim `0` and a dim before and after dim ``1``
    graph.output(ops.pad(x, [1, 1, 1, 1]))

    # [[0, 0, 0, 0], [0, 1, 2, 0], [0, 3, 4, 0], [0, 0, 0, 0]]

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

Parameters:

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

  • paddings (Iterable[int]) – The amount to pad. For a tensor of rank N, pass 2*N non-negative integers in the order [before_dim0, after_dim0, before_dim1, after_dim1, ...].

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

    How to fill the padded cells. Supported values:

    • "constant": fill using value.
    • "reflect": reflect the content across each edge, excluding the boundary element (like numpy.pad with mode='reflect').
    • "edge": repeat the nearest boundary element (like numpy.pad with mode='edge').

    Defaults to "constant".

  • value (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The fill value for mode="constant". Defaults to 0.

Returns:

A TensorValue representing the padded input, with the same dtype as input.

Raises:

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

Return type:

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.

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

device = DeviceRef.CPU()
with Graph("permute") as graph:
    # x has shape (1, 2, 3).
    x = ops.constant(
        [[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]],
        DType.float32,
        device=device,
    )
    # Reorder the dimensions to (2, 0, 1), producing shape (3, 1, 2).
    graph.output(ops.permute(x, [2, 0, 1]))

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

Parameters:

Returns:

A TensorValue representing x with its dimensions reordered to match dims. It has the same elements and dtype as x, with the order of the elements changed according to the permutation.

Raises:

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

Return type:

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 TensorValue with the broadcast shape representing 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

Computes the product of elements along a specified axis.

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

device = DeviceRef.CPU()
with Graph("prod_example") as graph:
    x = ops.constant(
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
        DType.float32,
        device=device,
    )
    graph.output(ops.prod(x, axis=-1))  # shape (2, 1): [[6.0], [120.0]]

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 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 computes the reduction along the last dimension. Defaults to -1.

Returns:

A TensorValue representing the product along axis. It has the same rank as x, with the axis dimension reduced to size 1.

Raises:

ValueError – If axis is out of range for the input’s rank.

Return type:

TensorValue

qmatmul()​

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

source

Performs matrix multiplication between floating point and quantized tensors.

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

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

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

Parameters:

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

Returns:

A TensorValue representing the dequantized, floating point result.

Raises:

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

Return type:

TensorValue

range()​

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

source

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

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

stop - start must be zero or have the same sign as step. Also, graph compilation fails when stop - start isn’t evenly divisible by step. For example, range(0, 5, 2) should produce three values, [0, 2, 4], but shape inference declares an output length of 2. The generated values therefore don’t fit the declared output shape.

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

device = DeviceRef.CPU()
with Graph("range_example") as graph:
    graph.output(ops.range(0, 5, 1, dtype=DType.float32, device=device))

model = InferenceSession().load(graph)
result = model.execute()[0]
# result holds [0.0, 1.0, 2.0, 3.0, 4.0].

Parameters:

Returns:

A TensorValue representing the generated sequence.

Raises:

  • ValueError – If out_dim is omitted for dynamic scalar inputs, if any input isn’t scalar, or if any input isn’t on the CPU.
  • RuntimeError – During graph compilation if a statically known interval isn’t evenly divisible by step, causing the inferred output length to disagree with the number of generated values.

Return type:

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

reduce_scatter_rms_norm()​

max.graph.ops.reduce_scatter_rms_norm(inputs, signal_buffers, gammas, epsilon, weight_offset=0.0)

source

Fused reduce-scatter sum + RMSNorm across devices (bf16 in/out).

Reduce-scatters inputs (one [rows, cols] tensor per device) along axis 0 across all devices, then RMSNorm-normalizes each device’s owned row shard in the same collective launch, keeping the reduced sum in float32 registers so there is no global-memory round-trip between the reduce-scatter and the norm. The norm is multiply_before_cast=True (gamma folded in float32, single cast to the input dtype last).

This is a full-world reduce-scatter (no device grouping): every input participates in one reduction. Rows are partitioned with the same ragged binning as reducescatter.sum() (remainder rows go to low ranks), so the sum output is a drop-in for reducescatter.sum along axis 0.

Parameters:

Returns:

A tuple (normed, residual) of two lists, each with one tensor per device: normed[i] is the RMSNorm of device i’s reduce-scatter shard and residual[i] is the reduce-scatter sum shard itself (the residual stream). Both have the input shape with axis 0 divided across devices.

Return type:

tuple[list[TensorValue], list[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), meaning negative values are set to zero while positive values are unchanged.

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 TensorValue of the same shape and dtype as x representing x with its negative elements 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 each element of a tensor along an axis.

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

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

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

device = DeviceRef.CPU()
with Graph("repeat_interleave_example") as graph:
    input = ops.constant(
        [[1.0, 2.0], [3.0, 4.0]], DType.float32, device=device
    )

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

    # Repeat row 0 twice and row 1 three times.
    repeats = ops.constant([2, 3], DType.int64, device=device)
    uneven_rows = ops.repeat_interleave(
        input, repeats=repeats, axis=0, out_dim=5
    )
    # [[1, 2], [1, 2], [3, 4], [3, 4], [3, 4]], shape (5, 2)
    graph.output(rows, uneven_rows)

Parameters:

  • x (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor.
  • repeats (int | TensorValue) – The number of times to repeat each element. Pass either a positive integer or a rank-0/rank-1 integer TensorValue.
  • axis (int | None) – The axis to repeat along. If None (the default), the input is flattened first.
  • out_dim (int | str | Dim | integer | TypedAttr | None) – The output size along axis. Required when repeats is a TensorValue.

Returns:

A TensorValue representing the input with its elements interleaved.

Raises:

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

Return type:

TensorValue

reshape()​

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

source

Reshapes a symbolic tensor.

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

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

device = DeviceRef.CPU()
with Graph("reshape") as graph:
    # x has shape (2, 3).
    x = ops.constant(
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
        DType.float32,
        device=device,
    )
    # Reshape the same 6 elements into shape (3, 2).
    graph.output(ops.reshape(x, [3, 2]))

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

Parameters:

Returns:

A TensorValue representing x with a new shape. The order and total number of elements stays the same as the input.

Raises:

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

Return type:

TensorValue

resize()​

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

source

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

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

device = DeviceRef.CPU()
with Graph("resize") as graph:
    # NCHW input: batch 1, 1 channel, 2x2 spatial.
    x = ops.constant(
        [[[[1.0, 2.0], [3.0, 4.0]]]], DType.float32, device=device
    )
    # Upscale the spatial dimensions to 4x4, shape (1, 1, 4, 4).
    graph.output(
        ops.resize(x, [1, 1, 4, 4], ops.InterpolationMode.BILINEAR)
    )

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

Parameters:

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

Returns:

A TensorValue representing the resized tensor with the given shape.

Raises:

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

Return type:

TensorValue

resize_bicubic()​

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

source

Resizes a tensor using bicubic interpolation.

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

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

device = DeviceRef.CPU()
with Graph("resize_bicubic") as graph:
    # NCHW input: batch 1, 1 channel, 2x2 spatial.
    x = ops.constant(
        [[[[1.0, 2.0], [3.0, 4.0]]]], DType.float32, device=device
    )
    # Upscale the spatial dimensions to 4x4, shape (1, 1, 4, 4).
    graph.output(ops.resize_bicubic(x, [1, 1, 4, 4]))

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

Parameters:

Returns:

A TensorValue representing the resized tensor, with shape size and the same dtype as input.

Raises:

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

Return type:

TensorValue

resize_linear()​

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

source

Resizes a tensor using linear (bilinear) interpolation.

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

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

device = DeviceRef.CPU()
with Graph("resize_linear") as graph:
    # NCHW input: batch 1, 1 channel, 2x2 spatial.
    x = ops.constant(
        [[[[1.0, 2.0], [3.0, 4.0]]]], DType.float32, device=device
    )
    # Upscale the spatial dimensions to 4x4, shape (1, 1, 4, 4).
    graph.output(ops.resize_linear(x, [1, 1, 4, 4]))

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

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]) – The desired output shape. Must have the same rank as input.

  • coordinate_transform_mode (int) –

    How to map an output coordinate to an input coordinate. Allowed values:

    • 0 (half_pixel): Default. Shifts by 0.5 before scaling, consistent with most deep learning frameworks.
    • 1 (align_corners): Aligns the corner pixels of the input and output so that the first and last coordinates are preserved exactly.
    • 2 (asymmetric): Applies no shift, mapping an output coordinate to out_coord / scale.
    • 3 (half_pixel_1D): Like half_pixel on every axis, except any axis whose output size is 1 maps to coordinate 0.

  • antialias (bool) – When True, applies an antialiasing filter when downscaling, which reduces aliasing artifacts by widening the tent filter support by 1 / scale. Has no effect when upscaling. Defaults to False.

Returns:

A TensorValue representing the resized tensor, with shape size and the same dtype as input.

Raises:

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

Return type:

TensorValue

resize_nearest()​

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

source

Resizes a tensor using nearest-neighbor interpolation.

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

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

device = DeviceRef.CPU()
with Graph("resize_nearest") as graph:
    # NCHW input: batch 1, 1 channel, 2x2 spatial.
    x = ops.constant(
        [[[[1.0, 2.0], [3.0, 4.0]]]], DType.float32, device=device
    )
    # Upscale the spatial dimensions to 4x4, shape (1, 1, 4, 4).
    graph.output(ops.resize_nearest(x, [1, 1, 4, 4]))

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

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]) – The desired output shape. Must have the same rank as input.

  • coordinate_transform_mode (int) –

    How to map an output coordinate to an input coordinate. Allowed values:

    • 0 (half_pixel). Default.
    • 1 (align_corners).
    • 2 (asymmetric).
    • 3 (half_pixel_1D).

    See resize_linear() for a description of each mode.

  • round_mode (int) –

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

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

Returns:

A TensorValue representing the resized tensor, with shape size and the same dtype as input.

Raises:

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

Return type:

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

Applies ROI-align pooling.

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

Parameters:

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

Returns:

A TensorValue representing the pooled values, with shape (num_rois, output_height, output_width, channels).

Raises:

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

Return type:

TensorValue

round()​

max.graph.ops.round(x)

source

Rounds a tensor to the nearest integer element-wise.

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

from max.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 TensorValue of the same shape and dtype as x representing each element of x 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 TensorValue of the same shape and dtype as x representing the reciprocal square root of each element of x.

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

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

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

device = DeviceRef.CPU()
with Graph("scatter") as graph:
    x = ops.constant([1, 2, 3, 4, 5], DType.int32, device=device)
    updates = ops.constant([10, 20], DType.int32, device=device)
    indices = ops.constant([0, 3], DType.int64, device=device)
    # Overwrite positions 0 and 3, producing [10, 2, 3, 20, 5].
    graph.output(ops.scatter(x, updates, indices, axis=0))

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

Parameters:

Returns:

A TensorValue representing input with updates written at indices. It has 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, if input, updates, and indices aren’t on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
  • Error – If input, updates, and indices don’t have equal rank, if updates.shape and indices.shape differ, or if any indices dimension exceeds the corresponding input dimension.

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 TensorValue representing the updated tensor. It has 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, if input, updates, and indices aren’t on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
  • Error – If input, updates, and indices don’t have equal rank, if updates.shape and indices.shape differ, or if any indices dimension exceeds the corresponding input dimension.

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 TensorValue representing the updated tensor. It has 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, if input, updates, and indices aren’t on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
  • Error – If input, updates, and indices don’t have equal rank, if updates.shape and indices.shape differ, or if any indices dimension exceeds the corresponding input dimension.

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 TensorValue representing the updated tensor. It has 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, if input, updates, and indices aren’t on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
  • Error – If input, updates, and indices don’t have equal rank, if updates.shape and indices.shape differ, or if any indices dimension exceeds the corresponding input dimension.

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 TensorValue representing the updated tensor. It has 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, if input, updates, and indices aren’t on the same device, or if any input is on a non-CPU device and strict_device_placement=DevicePlacementPolicy.Error.
  • Error – If input, updates, and indices don’t have equal rank, if updates.shape and indices.shape differ, or if any indices dimension exceeds the corresponding input dimension.

Return type:

TensorValue

scatter_nd()​

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

source

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

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

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

device = DeviceRef.CPU()
with Graph("scatter_nd") as graph:
    x = ops.constant(
        [[1, 2], [3, 4], [5, 6]], DType.int32, device=device
    )
    updates = ops.constant(
        [[10, 20], [50, 60]], DType.int32, device=device
    )
    indices = ops.constant([[0], [2]], DType.int64, device=device)
    # Overwrite rows 0 and 2, producing [[10, 20], [3, 4], [50, 60]].
    graph.output(ops.scatter_nd(x, updates, indices))

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

Parameters:

Returns:

A TensorValue representing input with updates scattered in. It has the same shape and dtype as input.

Raises:

ValueError – If the dtypes of input and updates mismatch, if indices dtype is not int32/int64, or if input, updates, and indices aren’t on the same device.

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 TensorValue representing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If the dtypes of input and updates mismatch, if indices dtype is not int32/int64, or if input, updates, and indices aren’t on the same device.

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 TensorValue representing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If the dtypes of input and updates mismatch, if indices dtype is not int32/int64, or if input, updates, and indices aren’t on the same device.

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 TensorValue representing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If the dtypes of input and updates mismatch, if indices dtype is not int32/int64, or if input, updates, and indices aren’t on the same device.

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 TensorValue representing the updated tensor. It has the same shape and dtype as input.

Raises:

ValueError – If the dtypes of input and updates mismatch, if indices dtype is not int32/int64, or if input, updates, and indices aren’t on the same device.

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

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. Must have a floating-point dtype.

Returns:

A TensorValue of the same shape and dtype as x representing each element of x mapped to the range (0, 1).

Raises:

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

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. Must have a floating-point dtype.

Returns:

A TensorValue of the same shape and dtype as x representing the SiLU activation applied to each element of x.

Raises:

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

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 TensorValue of the same shape and dtype as x representing the sine of each element of x.

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 of the input tensor based on indices.

The indices use NumPy-style slicing conventions, with one index per dimension. Each index is one of:

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

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

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

device = DeviceRef.CPU()
with Graph("slice_tensor") as graph:
    x = ops.constant(
        [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
        DType.int32,
        device=device,
    )
    # Take rows 0 and 1 and columns 1 and 2, producing
    # [[2, 3], [5, 6]].
    graph.output(ops.slice_tensor(x, [slice(0, 2), slice(1, 3)]))

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

Parameters:

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

Returns:

A TensorValue representing the sliced subtensor of x.

Raises:

  • IndexError – If an integer index or slice bound is out of range for its dimension.
  • ValueError – If x is a scalar, if an index TensorValue isn’t a scalar, if a slice step isn’t positive, or if an index form is unsupported.
  • TypeError – If an index for a symbolically-shaped dimension isn’t a slice or integer.
  • NotImplementedError – If a slice on a dynamic dimension omits the (slice, out_dim) form needed to compute its output size.

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 TensorValue of the same shape and dtype as value representing the softmax of value computed along axis.

Raises:

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

Return type:

TensorValue

split()​

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

source

Splits a tensor into several 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("split_example") as graph:
    x = ops.constant([1, 2, 3, 4, 5], DType.int32, device=device)

    # Split into a size-2 tensor and a size-3 tensor
    parts = ops.split(x, [2, 3], axis=0)  # Splits into [1, 2] and [3, 4, 5]
    graph.output(*parts)

model = InferenceSession().load(graph)
first, second = model.execute()

Parameters:

Returns:

A list of TensorValue objects, one per entry in split_sizes. Each has the same shape as x except along axis, where its size is the matching entry in split_sizes. If split_sizes is empty, returns an empty list.

Raises:

  • IndexError – If axis is out of range for the rank of x.
  • ValueError – If split_sizes doesn’t sum to the size of x along axis, or if any size is negative.

Return type:

list[TensorValue]

sqrt()​

max.graph.ops.sqrt(x)

source

Computes the 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("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]

Parameters:

Returns:

A TensorValue of the same shape and dtype as x representing the square root of each element of x.

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 dimension of size 1 from a symbolic tensor.

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

device = DeviceRef.CPU()
with Graph("squeeze") as graph:
    # x has shape (2, 1, 3).
    x = ops.constant(
        [[[1.0, 2.0, 3.0]], [[4.0, 5.0, 6.0]]],
        DType.float32,
        device=device,
    )
    # Remove the size-1 dimension at axis 1, producing shape (2, 3).
    graph.output(ops.squeeze(x, axis=1))

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 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, a value of -1 removes the last dimension.

Returns:

A TensorValue representing x with the dimension at axis removed. That dimension size must equal 1, so the result holds the same elements as x with one fewer dimension.

Raises:

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

Return type:

TensorValue

stack()​

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

source

Stacks tensors along a new axis.

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

device = DeviceRef.CPU()
with Graph("stack_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)

    # Stack the two (2, 2) tensors into one (2, 2, 2) tensor
    graph.output(ops.stack([a, b], axis=0))

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

Parameters:

  • values (Iterable[Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray]) – The tensors to stack. Each must have the same dtype, rank, and shape.
  • axis (int) – The position of the new axis. Negative values count from the end, where -1 inserts the new axis as the last dimension. Defaults to 0.

Returns:

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

Raises:

  • ValueError – If no tensors are provided, if the inputs don’t all have the same rank, if they don’t all have the same dtype, or if they aren’t all on the same device.
  • IndexError – If axis is out of range for the new tensor’s rank.

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 TensorValue representing 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

Computes the sum of elements along a specified axis.

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

device = DeviceRef.CPU()
with Graph("sum_example") as graph:
    x = ops.constant(
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
        DType.float32,
        device=device,
    )
    graph.output(ops.sum(x, axis=-1))  # shape (2, 1): [[6.0], [15.0]]

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 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 computes the reduction along the last dimension. Defaults to -1.

Returns:

A TensorValue representing the sum along axis. It has the same rank as x, with the axis dimension reduced to size 1.

Raises:

ValueError – If axis is out of range for the input’s rank.

Return type:

TensorValue

tanh()​

max.graph.ops.tanh(x)

source

Computes the hyperbolic tangent of a tensor element-wise.

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

from max.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 TensorValue of the same shape and dtype as x representing each element of x mapped to 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

Repeats a tensor along each of its dimensions.

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

This op runs on CPU. By default, an input on another device is copied to CPU for the operation (emitting a warning) and the result is copied back.

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

device = DeviceRef.CPU()
with Graph("tile_example") as graph:
    x = ops.constant([[1, 2], [3, 4]], DType.int32, device=device)

    # Repeat the columns twice, leaving the rows unchanged.
    graph.output(ops.tile(x, [1, 2]))  # [[1, 2, 1, 2], [3, 4, 3, 4]]

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

Parameters:

Returns:

A TensorValue representing the tiled input.

Raises:

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

Return type:

TensorValue

top_k()​

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

source

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

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

device = DeviceRef.CPU()
with Graph("top_k") as graph:
    x = ops.constant(
        [1.0, 3.0, 2.0, 5.0, 4.0], DType.float32, device=device
    )
    # The two largest values are 5 and 4 at indices 3 and 4.
    values, indices = ops.top_k(x, k=2, axis=-1)
    graph.output(values, indices)

model = InferenceSession().load(graph)
values, indices = model.execute()

Parameters:

  • input (Value[TensorType] | TensorValue | Shape | Dim | HasTensorValue | int | float | integer[Any] | floating[Any] | DLPackArray) – The input tensor from which to select the top k.
  • k (int) – The number of values to select from input. Must be in the range [0, input.shape[axis]].
  • axis (int) – The axis along which to select the top k. Defaults to -1.

Returns:

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

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.

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

device = DeviceRef.CPU()
with Graph("transpose") as graph:
    # x has shape (2, 3).
    x = ops.constant(
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
        DType.float32,
        device=device,
    )
    # Swap axes 0 and 1, producing shape (3, 2).
    graph.output(ops.transpose(x, 0, 1))

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 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, a value of -1 refers to the last axis.
  • axis_2 (int) – The other axis to transpose. If negative, this indexes from the end of the tensor.

Returns:

A TensorValue representing the input tensor with axis_1 and axis_2 transposed. It has the same elements and dtype as x, with the order of the elements changed according to the transposition. For a rank-zero tensor, axes -1 and 0 are accepted and the scalar is returned unchanged.

Raises:

IndexError – If axis_1 or axis_2 is out of range.

Return type:

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 TensorValue of the same shape and dtype as x representing each element of x truncated toward zero.

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 dimension of size 1 into a symbolic tensor.

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

device = DeviceRef.CPU()
with Graph("unsqueeze") as graph:
    # x has shape (3,).
    x = ops.constant(
        [1.0, 2.0, 3.0],
        DType.float32,
        device=device,
    )
    # Insert a size-1 dimension at axis 0, producing shape (1, 3).
    graph.output(ops.unsqueeze(x, axis=0))

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 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 to 1 plus the rank of the tensor. For example, a value of -1 adds a new dimension at the end, and -2 inserts the dimension immediately before the last dimension.

Returns:

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

Raises:

ValueError – If axis is out of bounds.

Return type:

TensorValue

where()​

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

source

Selects elements from x or y element-wise based on a condition.

At each position, takes the element from x where condition is true and the element from y where it’s false. The inputs are broadcast to a common shape. Either x or y can be a Python scalar, which is promoted to a tensor using the other operand’s dtype and device. At least one of x and y must be a tensor.

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

device = DeviceRef.CPU()
with Graph("where") as graph:
    condition = ops.constant(
        [True, False, True], DType.bool, device=device
    )
    x = ops.constant([1, 2, 3], DType.int32, device=device)
    y = ops.constant([10, 20, 30], DType.int32, device=device)
    # Take x where True and y where False, producing [1, 20, 3].
    graph.output(ops.where(condition, x, y))

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

Parameters:

Returns:

A TensorValue representing the element-wise selection from x and y according to condition.

Raises:

  • ValueError – If condition doesn’t have a boolean dtype, if condition, x, and y aren’t all on the same device, or if a Python scalar operand can’t be safely promoted to the other operand’s dtype.
  • TypeError – If neither x nor y is a tensor.
  • Error – If the shapes of condition, x, and y can’t be broadcast to a common shape.

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]

Models

Cookbook

Puzzles

Blog

Community

Contact