pytreeclass

Installation |Description |Quick Example |StatefulComputation |Benchamrks |Acknowledgements

🛠️ Installation

pip install pytreeclass

Install development version

pip install git+https://github.com/ASEM000/pytreeclass

📖 Description

pytreeclass is a JAX-compatible class builder to create and operate on stateful JAX PyTrees in a performant and intuitive way, by building on familiar concepts found in numpy, dataclasses, and others.

See documentation and 🍳 Common recipes to check if this library is a good fit for your work. If you find the package useful consider giving it a 🌟.

⏩ Quick Example

import jax
import jax.numpy as jnp
import pytreeclass as tc

@tc.autoinit
class Tree(tc.TreeClass):
    a: float = 1.0
    b: tuple[float, float] = (2.0, 3.0)
    c: jax.Array = jnp.array([4.0, 5.0, 6.0])

    def __call__(self, x):
        return self.a + self.b[0] + self.c + x


tree = Tree()
mask = jax.tree_map(lambda x: x > 5, tree)
tree = tree\
       .at["a"].set(100.0)\
       .at["b"][0].set(10.0)\
       .at[mask].set(100.0)

print(tree)
# Tree(a=100.0, b=(10.0, 3.0), c=[  4.   5. 100.])

print(tc.tree_diagram(tree))
# Tree
# ├── .a=100.0
# ├── .b:tuple
# │   ├── [0]=10.0
# │   └── [1]=3.0
# └── .c=f32[3](μ=36.33, σ=45.02, ∈[4.00,100.00])

print(tc.tree_summary(tree))
# ┌─────┬──────┬─────┬──────┐
# │Name │Type  │Count│Size  │
# ├─────┼──────┼─────┼──────┤
# │.a   │float │1    │      │
# ├─────┼──────┼─────┼──────┤
# │.b[0]│float │1    │      │
# ├─────┼──────┼─────┼──────┤
# │.b[1]│float │1    │      │
# ├─────┼──────┼─────┼──────┤
# │.c   │f32[3]│3    │12.00B│
# ├─────┼──────┼─────┼──────┤
# │Σ    │Tree  │6    │12.00B│
# └─────┴──────┴─────┴──────┘

# ** pass it to jax transformations **
# works with jit, grad, vmap, etc.

@jax.jit
@jax.grad
def sum_tree(tree: Tree, x):
    return sum(tree(x))

print(sum_tree(tree, 1.0))
# Tree(a=3.0, b=(3.0, 0.0), c=[1. 1. 1.])

📜 Stateful computations

Under jax.jit jax requires states to be explicit, this means that for any class instance; variables needs to be separated from the class and be passed explictly. However when using TreeClass no need to separate the instance variables ; instead the whole instance is passed as a state.

Using the following pattern,Updating state functionally can be achieved under jax.jit

import jax
import pytreeclass as tc

class Counter(tc.TreeClass):
    def __init__(self, calls: int = 0):
        self.calls = calls

    def increment(self):
        self.calls += 1
counter = Counter() # Counter(calls=0)

Here, we define the update function. Since the increment method mutate the internal state, thus we need to use the functional approach to update the state by using .at. To achieve this we can use .at[method_name].__call__(*args,**kwargs), this functional call will return the value of this call and a new model instance with the update state.

@jax.jit
def update(counter):
    value, new_counter = counter.at["increment"]()
    return new_counter

for i in range(10):
    counter = update(counter)

print(counter.calls) # 10

➕ Benchmarks

Benchmark flatten/unflatten compared to Flax and Equinox

CPU	GPU

Benchmark simple training against `flax` and `equinox`

Training simple sequential linear benchmark against flax and equinox

Num of layers	Flax/tc time	Equinox/tc time
10	1.427	6.671
100	1.1130	2.714

📙 Acknowledgements

Lenses
Treex, Equinox, tree-math, Flax PyTreeNode, TensorFlow, PyTorch
Lovely JAX

Metadata

Release files for pytreeclass 0.11.1

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for pytreeclass 0.11.1
File	Size	Uploaded
pytreeclass-0.11.1.tar.gz	60.1 kB	Aug 18, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for pytreeclass 0.11.1
File	Interpreter	ABI	Platform	Reset
pytreeclass-0.11.1-py3-none-any.whl 59.9 kB Aug 18, 2026	Python 3	none	any	Details

Total release size: 120.0 kB

Release files / pytreeclass-0.11.1.tar.gz

Download URL	pytreeclass-0.11.1.tar.gz
Size	60.1 kB
Tags	Source
SHA-256 checksum How to use checksums	`48cc9c700d1910fb58e660a15a17c9d53f5e028f483b41dbdd9ff5bdc7567325`
BLAKE2b-256 checksum How to use checksums	`b76e75072eb366c2a906bd731dd973db4489fba98da9262311b207377fee6712`
Upload date	Aug 18, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.14.7`

Release files / pytreeclass-0.11.1-py3-none-any.whl

Download URL	pytreeclass-0.11.1-py3-none-any.whl
Size	59.9 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`6724bb37d7c492a1d1b49fbc7d54d4029b3d751b57dc311d151120c60f9a81e5`
BLAKE2b-256 checksum How to use checksums	`f682ffb1720a42686954a89258873c97bad2789325159f4ddf05da8095ae802b`
Upload date	Aug 18, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.14.7`