Treepie documentation home

Treepie is a simple and powerful pytree creation and manipulation library intended to simplify JAX-programming. The intent of the library is to (a) let the programming experience be very close to “pure JAX-programming” (so no hardcoded or special-cased features relevant for specific applications like machine learning), but also (b) allow significant reduction in the boilerplate needed to implement complicated workflows (in some cases, even less boilerplate than in more specialized libraries).

Why bother with yet another JAX framework

Some assorted points that might pique your interest: With treepie, you can

  • implement stochastic layers like dropout without threading a random key (through the different functions leading up to and away from the stochastic layer),

  • implement stateful operations without threading a state object,

  • share weights (and more generally, let different branches of a pytree influence each other) without writing boilerplate code to ensure that only one reference to a shared weight crosses JAX-transform-boundaries (before getting copied to different locations),

  • add arbitrary tags and properties (e.g., nontrainable-tag or axis-names) to JAX arrays, without having to manually untag the arrays before using them.

Treepie doesn’t implement these features. Rather, it lets the user implement them easily. Also, treepie does this with

  • no special-casing or hardcoded features (pytrees created with treepie don’t need special processing. They can be processed/trained using other general purpose libraries in the JAX ecosystem like using equinox and optax),

  • no unsafe side-effects (so you can jax.jit away without worry),

  • no reliance on possibly unsupported JAX API, protocols, or features.

Yes way!

Getting started

Treepie is currently in early stages of development (say, pre-alpha). You can test it out as

pip install "treepie @ git+https://codeberg.org/prasanth/treepie.git"

Project repository

The source code for this project is hosted at https://codeberg.org/prasanth/treepie

Documentation

Documentation (work-in-progress) can be found at https://prasanth.codeberg.page/treepie