As an update to this functionality: @stef recently added the ability to provide entire MAX graphs as PyTorch custom operators. This significantly expands the scope of what can be used from MAX within a PyTorch custom operator to include our built-in highly optimized MAX kernels all the way to full graphs representing large portions of ML models.
To provide a MAX graph to a PyTorch custom operator, use the new @graph_op decorator around a function describing the MAX graph. An example of this in action is present in a new example in the modular repository.
We hope this enables an even easier path to start experimenting with MAX inside of a familiar PyTorch context, and allows for portions of a PyTorch model to be progressively replaced with optimized MAX operations and graphs.