Package: treex
Deprecation Notice: This library was an experiment trying to get pytree Modules working with Flax-like colletions. I'd currently recommend the following alternatives:Just custom pytrees:simple_pytreePytree module system:equinoxProduction ready module system:flaxTreexA Pytree Module system for Deep Learning in JAXMain Features💡Intuitive: Modules contain their own parameters and respect Object Oriented semantics like in PyTorch and Keras.🌳Pytree-based:  Modules are Pytrees whose leaves are its parameters, meaning they are fully compatible withjit,grad,vmap, etc.Treex is implemented on top ofTreeoand reexports all of its API for convenience.Getting Started|User Guide|Examples|DocumentationWhat is included?A baseModuleclass.Annmodule for with common layers implemented as wrappers over Flax layers.Alossesmodule with common loss functions.Ametricsmodule with common metrics.AnOptimizerclass that can wrap any optax optimizer.Why Treex?ShowDespite all JAX benefits, current Module systems are not intuitive to new users and add additional complexity not present in frameworks like PyTorch or Keras. Treex takes inspiration from S4TF and delivers an intuitive experience using JAX Pytree infrastructure.Current Alternative's Drawbacks and SolutionsCurrently we have many alternatives like Flax, Haiku, Objax, that have one or more of the following drawbacks:Module structure and parameter structure are separate, and parameters have to be manipulated around by the end-user, which is not intuitive. In Treex, parameters are stored in the modules themselves and can be accessed directly.Monadic architecture adds complexity. Flax and Haiku use anapplymethod to call modules that set a context with parameters, rng, and different metadata, which adds additional overhead to the API and creates an asymmetry in how Modules are being used inside and outside a context. In Treex, modules can be called directly.Among different frameworks, parameter surgery requires special consideration and is challenging to implement. Consider a standard workflow such as transfer learning, transferring parameters and state from a  pre-trained module or submodule as part of a new module; in different frameworks, we have to know precisely how to extract their parameters and how to insert them into the new parameter structure/dictionaries such that it is in agreement with the new module structure. In Treex, just as in PyTorch / Keras, we enable to pass the (sub)module to the new module, and parameters are automatically added to the new structure.Multiple frameworks deviate from JAX semantics and require particular versions ofjit,grad,vmap, etc., which makes it harder to integrate with other JAX libraries. Treex's Modules are plain old JAX PyTrees and are compatible with any JAX library that supports them.Other Pytree-based approaches like Parallax and Equinox do not have a total state management solution to handle complex states as encountered in Flax. Treex has the Filter and Update API, which is very expressive and can effectively handle systems with a complex state.InstallationInstall using pip:pipinstalltreexGetting StartedThis is a small appetizer to give you a feel for how using Treex looks like, be sure to checkout theUser Guidefor a more in-depth explanation.importtreexastximportnumpyasnpimportjax,optax# create some datax=np.random.uniform(size=(50,1))y=1.3*x**2-0.3+np.random.normal(size=x.shape)# initialize a Module, its simplemodel=tx.MLP([64,1]).init(key=42,inputs=x)# define an optimizer, init with model paramsoptimizer=tx.Optimizer(optax.adam(4e-3)).init(model)# define loss function, notice# Modules are jit-abel and differentiable 🤯@jax.jit@jax.graddefloss_fn(model:tx.MLP,x,y):# forward is a simple callpreds=model(x)# MSEreturn((preds-y)**2).mean()# basic training loopforstepinrange(500):# grads have the same type as modelgrads:tx.MLP=loss_fn(model,x,y)# apply the gradient updatesmodel=optimizer.update(grads,model)# Pytorch-like eval modemodel=model.eval()preds=model(x)Custom ModulesShowModules are TreeoTrees, which are Pytrees. When creating core layers you often mark fields that will contain state that JAX should be aware asnodesby assigning class variables to the output of functions liketx.Parameter.node():importtreexastxclassLinear(tx.Module):# use Treeo's API to define Parameter nodesw:jnp.ndarray=tx.Parameter.node()b:jnp.ndarray=tx.Parameter.node()def__init__(self,features_out:int):self.features_out=features_outdef__call__(self,x:jnp.ndarray)->jnp.ndarray:# init will call forward, we can know if we are inside itifself.initializing():# `next_key` only available during `init`key=tx.next_key()# leverage shape inferenceself.w=jax.random.uniform(key,shape=[x.shape[-1],self.features_out])self.b=jnp.zeros(shape=[self.features_out])# linear forwardreturnjnp.dot(x,self.w)+self.bmodel=Linear(10).init(key=42,inputs=x)Node field types (e.g.tx.Parameter) are called Kinds and Treex exports a whole family of Kinds which serve for differente purposes such as holding non-differentiable state (tx.BatchStats), metric's state (tx.MetricState), logging, etc. Checkout thekindssection for more information.Composite ModulesShowComposite Modules usually hold and call other Modules within them, while they would be instantiate inside__init__and used later in__call__like in Pytorch / Keras, in Treex you usually leverage the@tx.compactdecorator over the__call__method to define the submodules inline.classMLP(tx.Module):def__init__(self,features:Sequence[int]):self.features=features# compact lets you define submodules on the fly@tx.compactdef__call__(self,x:jnp.ndarray)->jnp.ndarray:forunitsinself.features[:-1]:x=Linear(units)(x)x=jax.nn.relu(x)returnLinear(self.features[-1])(x)model=MLP([32,10]).init(key=42,inputs=x)Under the hood all calls to submodule constructors (e.g.Linear(...)) insidecompactare assigned to fields in the parent Module (MLP) so they are part of the same Pytree, their field names are available under the._subtreesattribute.compactmust always define submodules in the same order.StatusTreex is in an early stage, things might break between versions but we will respect semanting versioning. Since Treex layers are numerically equivalent to Flax, it borrows some maturity and yields more confidence over its results. Feedback is much appreciated.Roadmap:Wrap all Flax Linen ModulesImplement more layers, losses, and metrics.Create applications and pretrained Modules.Contributions are welcomed!Sponsors 💚Quansight- paid development timeExamplesCheckout the/examplesdirectory for more detailed examples. Here are a few additional toy examples:Linear RegressionThis is a simple but realistic example of how Treex is used.fromfunctoolsimportpartialfromtypingimportUnionimportjaximportjax.numpyasjnpimportmatplotlib.pyplotaspltimportnumpyasnpimportoptaximporttreexastxx=np.random.uniform(size=(500,1))y=1.4*x-0.3+np.random.normal(scale=0.1,size=(500,1))# differentiate only w.r.t. parametersdefloss_fn(params,model,x,y):# merge params into modelmodel=model.merge(params)preds=model(x)loss=jnp.mean((preds-y)**2)# the model may contain state updates# so it should be returnedreturnloss,modelgrad_fn=jax.value_and_grad(loss_fn,has_aux=True)# both model and optimizer are jit-able@jax.jitdeftrain_step(model,x,y,optimizer):# select only the parametersparams=model.parameters()(loss,model),grads=grad_fn(params,model,x,y)# update params and modelparams=optimizer.update(grads,params)model=model.merge(params)# return new model and optimizerreturnloss,model,optimizermodel=tx.Linear(1).init(42,x)optimizer=tx.Optimizer(optax.adam(0.01)).init(model)forstepinrange(300):loss,model,optimizer=train_step(model,x,y,optimizer)ifstep%50==0:print(f"loss:{loss:.4f}")# eval mode "turns off" layers like Dropout / BatchNormmodel=model.eval()X_test=np.linspace(x.min(),x.max(),100)[:,None]preds=model(X_test)plt.scatter(x,y,c="k",label="data")plt.plot(X_test,preds,c="b",linewidth=2,label="prediction")plt.legend()plt.show()A Stateful ModuleHere is an example of creating a stateful module of aRollingMeanmetric and using them withjax.jit. For a real use cases use the metrics insidetreex.metrics.classRollingMean(tx.Module):count:jnp.ndarray=tx.State.node()total:jnp.ndarray=tx.State.node()def__init__(self):self.count=jnp.array(0,dtype=jnp.int32)self.total=jnp.array(0.0,dtype=jnp.float32)def__call__(self,x:jnp.ndarray)->jnp.ndarray:self.count+=np.prod(x.shape)self.total+=x.sum()returnself.total/self.count@jax.jitdefupdate(x:jnp.ndarray,metric:RollingMean)->Tuple[jnp.ndarray,RollingMean]:mean=metric(x)returnmean,metric# return mean value and updated metricmetric=RollingMean()foriinrange(10):x=np.random.uniform(-1,1,size=(100,1))mean,metric=update(x,metric)print(mean)