PAX documentation

PAX is a stateful pytree library for training neural networks using JAX. It is designed to be simple and easy to use while preserving benefits of JAX.

Installation

To install the latest version:

pip install git+https://github.com/ntt123/pax.git

Guides

• Getting started
• Train an image classifier
• Operators
• Understanding PAX’s module
• JAX transformations
• Improve performance

API Documentation

• PAX Basics
• Common Modules
• Module Transformations
• Utilities
• Experimental

PAX is licensed under the MIT License.

Indices

• Index