Skip to main content
Log in

Intro to custom ops

Custom operations (custom ops) extend MAX Graph's Python inference APIs with custom Mojo kernels. Whether you need to optimize performance of functions, implement custom algorithms, or create hardware-specific versions of existing operators, custom ops provide the flexibility you need.

The custom ops API provides complete control over MAX Graph while handling kernel integration and optimization pipelines automatically.

Try it now with our custom ops examples on GitHub or follow the Build custom ops for GPUs tutorial and let us know what you think.

How it works

A custom op consists of two main components that work together to integrate your custom implementation into the MAX execution pipeline:

  1. A custom function implementation written in Mojo that defines your computation
  2. A registration process that connects your function to the graph execution system

Under the hood, custom ops utilize high-level abstractions that handle memory management, device placement, and optimization. The graph compiler integrates your custom op implementation into the execution flow.

For more information: