Package: ninjax
🥷  Ninjax: Flexible Modules for JAXNinjax is a general module system forJAX. It gives the user complete
and transparent control over updating the state of each module, bringing the
flexibility of PyTorch and TensorFlow to JAX. Moreover, Ninjax makes it easy to
mix and match modules from different libraries, such asFlaxandHaiku.MotivationExisting deep learning libraries for JAX provide modules, but those modules
only specify neural networks and cannot easily implement training logic.
Orchestrating training all in one place, outside of the modules, is fine for
simple code bases. But it becomes a problem when there are many modules with
their own training logic and optimizers.Ninjax solves this problem by giving eachnj.Modulefull read and write
access to its state. This means modules can have train functions to
implement custom training logic, and call each other's train functions. Ninjax
is intended to be used with one or more neural network libraries, such asHaikuandFlax.The main differences to existing deep learning libraries are:Ninjax does not need separateapply()/init()functions. Instead, the
first function call creates variables automatically.Ninjax lets you access and update model parameters inside of impure
functions, so modules can handle their own optimizers and update logic.Natural support for modules with multiple functions without need for
Flax'ssetup()function or Haiku'shk.multi_transform().Ninjax' flexible state handling makes it trivial to mix and match
modules from other deep learning libraries in your models.InstallationNinjax isa single file, so you can just copy it to your project
directory. Or you can install the package:pip install ninjaxQuickstartimportflaximportjaximportjax.numpyasjnpimportninjaxasnjimportoptax# Supports all Flax and Haiku modulesLinear=nj.FromFlax(flax.linen.Dense)classMyModel(nj.Module):lr:float=0.01def__init__(self,size):self.size=sizeself.opt=optax.adam(self.lr)# Define submodules upfrontself.h1=Linear(128,name='h1')self.h2=Linear(128,name='h2')defpredict(self,x):x=jax.nn.relu(self.h1(x))x=jax.nn.relu(self.h2(x))# Define submodules inlinex=self.get('h3',Linear,self.size,use_bias=False)(x)# Create state entries inlinex+=self.get('bias',jnp.zeros,self.size)returnxdeftrain(self,x,y):# Gradient with respect to submodules or state entrieskeys=[self.h1,self.h2,f'{self.path}/h3',f'{self.path}/bias']loss,params,grads=nj.grad(self.loss,keys)(x,y)# Update weightsoptstate=self.get('optstate',self.opt.init,params)updates,optstate=self.opt.update(grads,optstate)new_params=optax.apply_updates(params,updates)self.put(new_params)# Store the new paramsreturnlossdefloss(self,x,y):return((self.predict(x)-y)**2).mean()# Create model and example datamodel=MyModel(3,lr=0.01,name='model')x=jnp.ones((64,32),jnp.float32)y=jnp.ones((64,3),jnp.float32)# Populate initial state from one or more functionsstate={}state=nj.init(model.train)(state,x,y,seed=0)print(state['model/bias'])# [0., 0., 0.]# Purify and jit one or more functionstrain=nj.pure(model.train)train=jax.jit(train)# Training loopforx,yin[(x,y)]*10:state,loss=train(state,x,y)print('Loss:',float(loss))# Look at the parametersprint(state['model/bias'])# [-1.2e-09  1.8e-08 -2.5e-09]TutorialHow can I create state entries?Ninjax gives modules full control over reading and updating their state
entries. Useself.get(name, ctor, *args, **kwargs)to define state entries.
The first call creates the entry asctor(*args, **kwargs). Later calls return
the current value:classModule(nj.Module):defcompute(self,x):init=jax.nn.initializers.variance_scaling(1,'fan_avg','uniform')weights=self.get('weights',init,nj.rng(),(64,32))bias=self.get('bias',jnp.zeros,(32,),jnp.float32)print(self.getm())# {'/path/to/module/weights': ..., '/path/to/module/bias': ...}returnx@weights+biasHow can I update state entries?To update the state entries of a module, useself.put(name, value)for
individual entries ofself.putm(mapping)to update multiple values:classModule(nj.Module):defcounting(self):counter=nj.get('counter',jnp.zeros,(),jnp.int32)self.put('counter',counter+1)print(self.get('counter'))# 1state=self.getm()state['counter']+=1self.putm(state)print(self.getm()['counter'])# 2print(self.get('counter'))# 2How can I use JIT compilation?Thenj.pure()function makes the state your JAX code uses explicit, so it can
be transformed freely:model=Model()train=jax.jit(nj.pure(model.train))params={}result,params=train(param,rng,...)Calling the pure function will create any parameters that are not yet in theparamsdictionary and return the new parameters alongside the function
output.You can speed up the first function call (where parameters are created) by
usingnj.jitinstead ofjax.jit. Internally, this avoids compiling two
versions of the function.How can I compute gradients?You can usejax.gradas normal for computing gradients with respect to
explicit inputs of your function. To compute gradients with respect to Ninjax
state, usenj.grad(fn, keys):classModule(nj.Module):deftrain(self,x,y):params=self.getm('.*')loss,grads=nj.grad(self.loss,params.keys())(x,y)params=jax.tree_map(lambdap,g:p-0.01*g,params,grads)self.putm(params)Theself.getm(filter='.*')method optionally accepts a regex pattern to select
only a subset of the state dictionary. It also returns only state entries of
the current module. To access the global state, usenj.state().How can I define modules compactly?You can useself.get(name, ctor, *args, **kwargs)inside methods of your
modules. When called for the first time, it creates a new state entry from the
constructorctor(*args, **kwargs). Later calls return the existing entry:classModule(nj.Module):def__call__(self,x):x=jax.nn.relu(self.get('h1',Linear,128)(x))x=jax.nn.relu(self.get('h2',Linear,128)(x))x=self.get('h3',Linear,32)(x)returnxHow can I use Haiku modules?There is nothing special about using external libraries with Ninjax. Haiku
requires its modules to be passed throughhk.transformand the initialized
viatransformed.init(rng, batch). For convenience, Ninjax providesnj.HaikuModuleto do this for you:classModule(nj.Module):def__init__(self):self.mlp=nj.HaikuModule(hk.nets.MLP,[128,128,32])def__call__(self,x):returnself.mlp(x)You can also predefine a list of aliases for Haiku modules that you want to use
frequently:Linear=functools.partial(nj.HaikuModule,hk.Linear)Conv2D=functools.partial(nj.HaikuModule,hk.Conv2D)MLP=functools.partial(nj.HaikuModule,hk.nets.MLP)# ...How can I use Flax modules?There is nothing special about using external libraries with Ninjax. Flax
requires its modules to be initialized viaparams = model.init(rng, batch)and used viamodel.apply(params, data). For convenience, Ninjax providesnj.FlaxModuleto do this for you:classModule(nj.Module):def__init__(self):self.linear=nj.FlaxModule(nn.Dense,128)def__call__(self,x):returnself.linear(x)You can also predefine a list of aliases for Flax modules that you want to use
frequently:Dense=functools.partial(nj.FlaxModule,nn.Dense)Conv=functools.partial(nj.FlaxModule,nn.Conv)# ...How can I use Optax optimizers?There is nothing special about using external libraries like Optax with Ninjax.
Optax requires its optimizers to be initialized, their state to be passed
through the optimizer call, and the resulting updates to be applied. For
convenience, Ninjax providesnj.OptaxModuleto do this for you:classModule(nj.Module):def__init__(self):self.mlp=MLP()self.opt=nj.OptaxModule(optax.adam,1e-3)deftrain(self,x,y):self.mlp(x)# Ensure paramters are created.metrics=self.opt(self.mlp.getm('.*'),self.loss,x,y)returnmetrics# {'loss': ..., 'grad_norm': ...}defloss(self,x,y):return((self.mlp(x)-y)**2).mean()QuestionsIf you have a question, pleasefile an issue.