Visualize, create, and operate on JAX PyTree in the most intuitive way possible.

Installation

Install from pip:

pip install pytreeclass

User guides

• 🏃 Getting started
• 📜 Stateful computations
• 🥶 Dealing with non-jax types
• 🍳 Common recipes

Examples

• 🧮 Building a mini optimizer library

API Documentation

• pytreeclass API

Apache2.0 License.

Indices

• Index