refrax 0.0.4

Chainable optics for JAX PyTrees

Refrax

Refrax is a tiny library implementing the optics functional pattern for JAX PyTrees.

Installation

Refrax can be installed using pip:

pip install refrax

Quick example

This example demonstrates the Mixin approach on an Equinox module. This simply adds .at to your class, which returns refrax.Lens(self).

We define a dummy model with the mixin:

from refrax import OpticsMixin

import equinox as eqx
import jax

class Model(eqx.Module, OpticsMixin):
    core: eqx.nn.Linear
    head: eqx.nn.Linear
    dropout: float

key1, key2 = jax.random.split(jax.random.key(0))
model = Model(
    core=eqx.nn.Linear(in_features=5, out_features=5, key=key1),
    head=eqx.nn.Linear(in_features=5, out_features=2, key=key2),
    dropout=0.5
)

Then we can do updates using .at:

model = model.at.dropout.set(0.1)
model = model.at.select("core", "head").bias.apply(lambda b: b + 1.0)

model.dropout
# 0.1

model.core.bias
# [0.646, 0.860, 0.670 , 1.277 , 0.727]

model.head.bias
# [1.634, 1.877]

Documentation

Documentation is available here.

Related

The library uses Equinox (specifically eqx.tree_at) under the hood.