Skip to content

Latest commit

 

History

History
 
 

README.md

MAX examples

These examples demonstrate the power and flexibility of MAX. They include:

A collection of sample programs written in the Mojo programming language.

The MAX Graph API provides a powerful framework for staging computational graphs to be run on GPUs, CPUs, and more. Each operation in one of these graphs is defined in Mojo, an easy-to-use language for writing high-performance code.

The examples here illustrate how to construct custom graph operations in Mojo that run on GPUs and CPUs, as well as how to build computational graphs that contain and run them on different hardware architectures.

In addition to placing custom Mojo functions within a computational graph, the MAX Driver API can handle direct compilation of GPU functions written in Mojo and can dispatch them onto the GPU. This is a programming model that may be familiar to those who have worked with CUDA or similar GPGPU frameworks.

These examples show how to compile and run Mojo functions, from simple to complex, on an available GPU. Note that a MAX-compatible GPU will be necessary to build and run these.

To enable progressive introduction of Mojo into an existing Python codebase, Mojo modules and functions can be referenced as if they were native Python code. This interoperability between Python and Mojo can allow for slower Python algorithms to be selectively replaced with faster Mojo alternatives.

These examples illustrate how that can work, including using Mojo functions running on a compatible GPU.

PyTorch custom operations can be defined in Mojo to try out new algorithms on GPUs. These examples show how to extend PyTorch layers using custom operations written in Mojo.

MAX has the power to accelerate existing PyTorch models directly, and provides Python, Mojo, and C APIs for this. These examples showcase common models and how to run them even faster via MAX.

Jupyter notebooks that showcase PyTorch models being accelerated through MAX.

The MAX Python API provides a PyTorch-like interface for building neural network components that compile to highly optimized graphs. These examples demonstrate how to create reusable, modular components using MAX's nn.Module class.

The examples include custom layers, blocks, and architectural patterns that showcase the flexibility of MAX's Python API for deep learning development, from simple MLP blocks to more complex neural network architectures.