Skip to main content
Max PlatformDocsCompanyBlogModCon
Max PlatformDocsCompanyBlogModCon
Contact salesSign in

Intro to MAX extensibility

The AI model you get from a framework like PyTorch or TensorFlow is built as a graph of connected operations ("ops"). Although most ops are simple math functions, efficiently executing models that include trillions of ops requires highly-performant implementations for each op (sometimes called "kernels"). However, even the most high-performance ops aren't enough to achieve peak performance. It's also necessary to employ graph compilers that can analyze the entire graph and optimize the calculations and memory that span across a sequence of ops.

That's why MAX Engine is designed to be fully extensible with Mojo. Regardless of the model format you have (such as PyTorch, ONNX, or MAX Graph), you can write custom ops in Mojo that the MAX Engine compiler can natively analyze and optimize along with with the rest of the model.

There’s no other AI platform that provides graph extensibility that’s as simple and powerful as this. All you do is write your custom op in Mojo (using code that reads like Python), package it with the mojo package command, and pass it to MAX Engine when you load your model.

  • MAX will parse, analyze, optimize, and fuse your op with the rest of the model. (Fusing is coming soon.)
  • Custom ops are platform independent.
  • Custom ops work with any model format supported by MAX.
  • Custom ops are optimized exactly like our ops; they are not treated as second-class citizens in the MAX Engine compiler.

Not only does MAX make it easy to load custom ops for any model, but by writing them in Mojo, you're working in a Python-friendly language that provides high-performance systems programming features (including strong type checking, memory safety, hardware intrinsics in MLIR/LLVM, and even inline assembly).

Overview​

The process to add a custom op to your model includes the following steps:

  1. Implement your op as a Mojo function that takes a Tensor for each op input, and returns a Tensor as its output.
    1. Specify the name of your op with the max.register.op() function decorator.
    2. Package your op with the mojo package command.
  2. Add the op to your model.
    1. If you’re loading a PyTorch or ONNX model, specify the path to your custom op package with the custom_ops_path argument when you call load().
    2. If you’re building your model with MAX Graph, add the custom op to the graph with ops.custom().

For a step-by-step walkthrough, see the following pages: