Package: treeo
TreeoA small library for creating and manipulating custom JAX Pytree classesLight-weight: has no dependencies other thanjax.Compatible: TreeoTreeobjects are compatible with anyjaxfunction that accepts Pytrees.Standards-based:treeo.fieldis built on top of python'sdataclasses.field.Flexible: Treeo is compatible with both dataclass and non-dataclass classes.Treeo lets you easily create class-based Pytrees so your custom objects can easily interact seamlessly with JAX. Uses of Treeo can range from just creating simple simple JAX-aware utility classes to using it as the core abstraction for full-blown frameworks. Treeo was originally extracted from the core ofTreexand shares a lot in common withflax.struct.Documentation|User GuideInstallationInstall using pip:pipinstalltreeoBasicsWith Treeo you can easily define your own custom Pytree classes by inheriting from Treeo'sTreeclass and using thefieldfunction to declare which fields are nodes (children) and which are static (metadata):importtreeoasto@dataclassclassPerson(to.Tree):height:jnp.array=to.field(node=True)# I am a node field!name:str=to.field(node=False)# I am a static field!fieldis just a wrapper arounddataclasses.fieldso you can define your Pytrees as dataclasses, but Treeo fully supports non-dataclass classes as well. Since allTreeinstances are Pytree they work with the various functions from thejaxlibrary as expected:p=Person(height=jnp.array(1.8),name="John")# Trees can be jitted!jax.jit(lambdaperson:person)(p)# Person(height=array(1.8), name='John')# Trees can be mapped!jax.tree_map(lambdax:2*x,p)# Person(height=array(3.6), name='John')KindsTreeo also include a kind system that lets you give semantic meaning to fields (what a field represents within your application). A kind is just a type you pass tofieldvia itskindargument:classParameter:passclassBatchStat:passclassBatchNorm(to.Tree):scale:jnp.ndarray=to.field(node=True,kind=Parameter)mean:jnp.ndarray=to.field(node=True,kind=BatchStat)Kinds are very useful as a filtering mechanism viatreeo.filter:model=BatchNorm(...)# select only Parameters, mean is filtered outparams=to.filter(model,Parameter)# BatchNorm(scale=array(...), mean=Nothing)Nothingbehaves likeNonein Python, but it is a special value that is used to represent the absence of a value within Treeo.Treeo also offers themergefunction which lets you rejoin filtered Trees with a logic similar to Pythondict.updatebut done recursively:defloss_fn(params,model,...):# add traced params to modelmodel=to.merge(model,params)...# gradient only w.r.t. paramsparams=to.filter(model,Parameter)# BatchNorm(scale=array(...), mean=Nothing)grads=jax.grad(loss_fn)(params,model,...)For a more in-depth tour check out theUser Guide.ExamplesA simple Treefromdataclassesimportdataclassimporttreeoasto@dataclassclassCharacter(to.Tree):position:jnp.ndarray=to.field(node=True)# node fieldname:str=to.field(node=False,opaque=True)# static fieldcharacter=Character(position=jnp.array([0,0]),name='Adam')# character can freely pass through jit@jax.jitdefupdate(character:Character,velocity,dt)->Character:character.position+=velocity*dtreturncharactercharacter=update(charactervelocity=jnp.array([1.0,0.2]),dt=0.1)A Stateful Treefromdataclassesimportdataclassimporttreeoasto@dataclassclassCounter(to.Tree):n:jnp.array=to.field(default=jnp.array(0),node=True)# nodestep:int=to.field(default=1,node=False)# staticdefinc(self):self.n+=self.stepcounter=Counter(step=2)# Counter(n=jnp.array(0), step=2)@jax.jitdefupdate(counter:Counter):counter.inc()returncountercounter=update(counter)# Counter(n=jnp.array(2), step=2)# map over the treeFull Example - Linear Regressionimportjaximportjax.numpyasjnpimportmatplotlib.pyplotaspltimportnumpyasnpimporttreeoastoclassLinear(to.Tree):w:jnp.ndarray=to.node()b:jnp.ndarray=to.node()def__init__(self,din,dout,key):self.w=jax.random.uniform(key,shape=(din,dout))self.b=jnp.zeros(shape=(dout,))def__call__(self,x):returnjnp.dot(x,self.w)+self.b@jax.value_and_graddefloss_fn(model,x,y):y_pred=model(x)loss=jnp.mean((y_pred-y)**2)returnlossdefsgd(param,grad):returnparam-0.1*grad@jax.jitdeftrain_step(model,x,y):loss,grads=loss_fn(model,x,y)model=jax.tree_map(sgd,model,grads)returnloss,modelx=np.random.uniform(size=(500,1))y=1.4*x-0.3+np.random.normal(scale=0.1,size=(500,1))key=jax.random.PRNGKey(0)model=Linear(1,1,key=key)forstepinrange(1000):loss,model=train_step(model,x,y)ifstep%100==0:print(f"loss:{loss:.4f}")X_test=np.linspace(x.min(),x.max(),100)[:,None]y_pred=model(X_test)plt.scatter(x,y,c="k",label="data")plt.plot(X_test,y_pred,c="b",linewidth=2,label="prediction")plt.legend()plt.show()