Skip to main content

custom_ops

Helpers for custom ops.

custom

custom[name: StringLiteral](owned values: SymbolTuple, owned out_type: AnyMOType) -> Symbol

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

The custom op should be registered by annotating a function with the max.register.op decorator. That function should have exactly one output.

Parameters:

  • name (StringLiteral): The op name provided to max.register.op.

Args:

  • values (SymbolTuple): The inputs to the op function.
  • out_type (AnyMOType): The type of the single output to the function.

Returns:

A symbolic value representing the output of the op in the graph.

custom_nv

custom_nv[name: StringLiteral](owned values: SymbolTuple, owned out_types: TypeTuple) -> SymbolTuple

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

The custom op should be registered by annotating a function with the max.register.op decorator. That function should have exactly one output.

Parameters:

  • name (StringLiteral): The op name provided to max.register.op.

Args:

  • values (SymbolTuple): The inputs to the op function.
  • out_types (TypeTuple): The types of each output to the function.

Returns:

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