For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

CustomOp

CustomOp

class max.experimental.CustomOp(name, inputs, outputs, extensions, parameters, declared_inputs, declared_outputs, pinned_params)

source

Bases: object

A declared custom op; call it eagerly or in a graph.

Build one with declare(), which validates the signature and compiles the kernel packages once. Instances are immutable and compare by declaration.

The same op accepts an eager Tensor, running the kernel immediately, or a TensorValue inside a Graph build, staging a mo.custom op. Output types come from the declared signature, so nothing is computed per call beyond binding the signature’s symbols to the operands’ dims.

from max.dtype import DType
from max.experimental import Tensor, custom
from max.graph import DeviceRef, Graph, TensorType

n, w, c = custom.Symbols("n", "w", "c")
downsample = custom.declare(
    "downsample",
    inputs={"x": custom.TemplateType(DType.float32, [n, w, c])},
    outputs=[custom.TemplateType(DType.float32, [n, w // 2, c])],
    custom_extensions=["path/to/kernels"],
)

y = downsample(Tensor.ones((2, 8, 3)))

# Graph: staged, output type `[n, w // 2, c]` -> `["n", 4, "c"]`.
x_type = TensorType(DType.float32, ["n", 8, "c"], DeviceRef.CPU())
with Graph("g", input_types=[x_type]) as graph:
    graph.output(downsample(graph.inputs[0]))

declared_inputs

declared_inputs: Mapping[str, TemplateType]

source

The signature as declared, before valued parameters were folded into inputs/outputs, so __getitem__() can re-specialize a Param. Derived from the fields above, so not compared.

declared_outputs

declared_outputs: tuple[TemplateType, ...]

source

extensions

extensions: tuple[Path, ...]

source

Kernel packages, already compiled to binaries.

inputs

inputs: Mapping[str, TemplateType]

source

One entry per operand, keyed by operand name, in operand order. A dict is not hashable, so the hash covers the other fields; equality still covers all of them.

is_complete

property is_complete: bool

source

Whether the signature has no unresolved Param left.

False means at least one parameter a dim names, or one declared bare, has no value, so the def cannot be staged (see _stage()) until it is given one, either at declaration or through __getitem__(). A parameter whose value could never fold into a dim is refused at declaration (see _check_dim_param_value()) rather than stranding a def here with no remedy.

name

name: str

source

The registered kernel symbol.

outputs

outputs: tuple[TemplateType, ...]

source

One entry per result.

parameters

parameters: Mapping[str, bool | int | float | str | DType | None]

source

Compile-time parameters passed to the kernel. None declares one without giving it a value yet (see is_complete).

pinned_params

pinned_params: Mapping[tuple[str, int], str]

source

Input dims a Param pinned before folding, by (operand, dim index), so a static-mismatch error in _unify() can still name the parameter – by then the dim is a bare static.

token

property token: str

source

The registry token for this op, a digest of its declaration.

Content-derived, so identical declarations share one token and so one compile: the token is stamped into a discardable attribute, which prints into the module ASM that executor._eager_model_cache_key hashes.

Covers exactly what a resolved op is read for: kernel symbol, parameters, resolved signature, and the extension paths custom_gc.binding_for links against. Kernel library CONTENT is deliberately absent: both consumers already hash it themselves – BindingKey.lib_hashes on the interpreter path, and _eager_model_cache_key’s resolved-library component alongside the ASM hash on the compile path – so a kernel edit is covered without reading every library file at every declaration.