Package: nnx
NNXNeuralNetworks for JAXNNX is a Neural Networks library for JAX that provides a simple yet powerful module system that adheres to standard Python semantics. Its aim is to combine the robustness ofFlaxwith a simplified, Pythonic API akin to that ofPyTorch.Pythonic: Modules are just regular python classes, they contain their own state, are fully mutable, and allow sharing references between Modules.Compatible: Easily convert back and forth between Modules and pytrees using the Functional API to integrate with any JAX API.Safe: NNX incorporates mechanisms to try to prevent tracer leakage, avoid stale RNGs, and ensure proper state propagation in order to help produce correct JAX programs.Semantic: Partition a Module's state into different semantic collections, allowing for fine-grained control when applying JAX transformations.Table of ContentsInstallationGetting StartedFAQsExamplesUser GuideInstallationTo get started withnnx, install the package via pip:pip install nnxFor the most recent version, install directly from our GitHub repository:pip install git+https://github.com/cgarciae/nnxGetting StartedThe following example guides you through creating a basicLinearmodel with NNX and executing a forward pass. It also demonstrate how handle mutable state by showing how to keep track of the number of times the model has been called.importnnximportjaximportjax.numpyasjnpclassCount(nnx.Variable):passclassLinear(nnx.Module):def__init__(self,din:int,dout:int,*,ctx:nnx.Context):key=ctx.make_rng("params")self.w=nnx.Param(jax.random.uniform(key,(din,dout)))self.b=nnx.Param(jnp.zeros((dout,)))self.count=Count(0)# track the number of callsdef__call__(self,x):self.count+=1returnx@self.w+self.bmodel=Linear(din=12,dout=2,ctx=nnx.context(0))# Forward pass and verify the call countx=jnp.ones((8,12))y=model(x)assertmodel.count==1In this examplennx.context(0)create aPRNGKeyforparamswith seed0, this is used bymake_rnginside__init__to generate a random key to initialize the parameters.Training with the Functional APITheFunctional APIconverts an NNX Module python semantics into pure pytree object with functional semantics. It is the recommended way to use NNX as it provides tight control over the state, allows you to use regular JAX transformations, and it minimizes overhead. In this example the model will be trained using Stochastic Gradient Descent (SGD).(params,counts),moduledef=model.partition(nnx.Param,Count)@jax.jitdeftrain_step(params,counts,x,y):defloss_fn(params):y_pred,(updates,_)=moduledef.apply(params,counts)(x)loss=jax.numpy.mean((y_pred-y)**2)returnloss,updates.filter(Count)# compute gradientgrads,counts=jax.grad(loss_fn,has_aux=True)(params)# SGD updateparams=jax.tree_map(lambdaw,g:w-0.1*g,params,grads)returnparams,counts# execute the training stepparams,counts=train_step(params,counts,x,y)model=moduledef.merge(params,counts)assertmodel.count==2Training with Lifted TransformsLifted Transformsprovide a convenient way interact with NNX Modules. In this example, we use thennx.jitandnnx.gradlifted transforms to define the training step. The model is trained using Stochastic Gradient Descent (SGD). Because lifted transforms automatically update the Module's state,train_stepdoesn't require a return statement.@nnx.jitdeftrain_step(model,x,y):defloss_fn(model):y_pred=model(x)returnjax.numpy.mean((y_pred-y)**2)# compute gradientgrads:nnx.State=nnx.grad(loss_fn,wrt=nnx.Param)(model)# SGD updatemodel.update_state(jax.tree_map(lambdaw,g:w-0.1*g,model.filter(nnx.Param),grads))# execute the training steptrain_step(model,x,y)assertmodel.count==2Note: Usingnnx.jitintroduces some overhead when compared to usingjax.jitdirectly. Usennx.jitfor simple prototypes, but for production code usejax.jitdirectly.ExamplesUsing the Functional API: Shows how to train a simple model using the functional API.Using Lifted Transforms: Shows how to train a simple model using lifted transforms.Using TrainState: Shows how to train a simple model using the functional API with the help ofTrainState.Using PureModule(experimental): Shows how to train a simple model using the functional API and leveragingPureModuleto simplify the code.Training a VAE: Shows how to train a VAE on the binarized MNIST dataset, uses the functional API,TrainState, and shows how to use capture intermediate values to retrievekl_loss.Scan over layers: An contrived example that implements scan over layers with dropout and a share BatcNorm layer to showcase how lifted transforms can be implemented. It uses the functional API along withjax.vmapandjax.lax.scan.Creating a Transformer: Shows how to create a Transformer with an auto-regressive decoder that uses scan over layers and a kv-cache for fast inference. Credits to @levskaya.FAQsStatusNNX is still in early development so expect bugs and breaking changes. That said, current API is the result of months of experimentation and we don't expect any major changes in the near future.How is it different from Flax?NNX takes the best features that allow Flax to scale to large projects and integrates them into a much simpler Module system with pythonic semantics.One place in which NNX strongly deviates from Flax is that (currently) it avoids shape inference in favor of static initialization. It is not a technical limitation but rather a design choice. This design both simplifies the internal implementation and makes it easier to reason about the code for the user, at the cost of being more verbose at times. On the other hand, Pytorch users will feel right at home.How is it different from Equinox?While they might look similar at a surface-level, NNX's Module system is more powerful and flexible than Equinox's, it contains the following additional features:Uses regular python classes (no mandatory dataclass behavior).Modules are mutableReference sharing between Modules is allowedMutable state lives inside the Module (no need for a separateState container).Supports node metadata and semantic partitioning.One major difference between the two frameworks is that, by design, NNX Modules are not Pytrees. This adds a safety layer as it prevents state updates from being lost by accident due to referential transparency. It also removes the need of threading a separateState containerthroughout the code in order to propagate state. In NNX state updates are either always preserved or explicitly discarded by the user.User GuideModulesNNX Modules are normal python classes, they obey regular python semantics such as mutability and reference sharing, including reference cycles. They can contain 2 types of attributes: node attributes and static attributes. Node attributes include NNXVariables (e.g.nnx.Param), Numpy arrays, JAX arrays, submodules Modules, and other NNX types. All other types are treated as static attributes.classFoo(nnx.Module):def__init__(self,ctx:nnx.Context):# node attributesself.variable=nnx.Param(jnp.array(1))self.np_buffer=np.array(2)self.jax_buffer=jnp.array(3)self.node=nnx.Node([4,5])self.submodule=nnx.Linear(2,4,ctx=ctx)# static attributesself.int=1self.float=2.0self.str="hello"self.list=[1,2,3]model=Foo(din=12,dout=2,ctx=nnx.context(0))As shown above, python container types such aslist,tuple, anddictare treated as static attributes, if similar functionality is needed, NNX provides theSequenceandDictModules.Functional APINNX Modules are not pytrees so they cannot be passed to JAX transformations. In order to interact with JAX, a Module must be partitioned into aStateandModuleDefobjects. TheStateobject is a flat dictionary-like pytree structure that contains all the deduplicated node attributes, and theModuleDefcontains the static attributes and structural information needed to reconstruct the Module.state,moduledef=model.partition()State({
  ('jax_buffer',): Array(3),
  ('node',): Node(value=[4, 5]),
  ('np_buffer',): array(2),
  ('submodule', 'bias'): Param(value=Array(...)),
  ('submodule', 'kernel'): Param(value=Array(...)),
  ('variable',): Param(value=Array(1))
})StateandModuleDefare pytrees so they can be passed to JAX transformations. More over,ModuleDefprovides 2 very important methods:mergeandapply. Themergemethod can be used to create a newModulefrom aStateobject:model=moduledef.merge(state)This can be use to e.g. recreate a module inside a JAX transformation. Theapplyprovides a functional interface to the module, it can be used call any method or submodule and get the output and the updated state:# run __call__y,(state,moduledef)=moduledef.apply(state)(x)# run some_methody,(state,moduledef)=moduledef.apply(state).some_method(x)# run submoduley,(state,moduledef)=moduledef.apply(state).submodule(x)applycan call any nested method or submodule as long as it can be accessed via the.or[]operators.Partitioning StateIn NNX you can filter based on any node type, most commonly you will want to filter based onnnx.Variablesubclasses such asnnx.Paramornnx.BatchStat.Here are various examples of how you can use thepartitionmethod to split a module into multiple substates:# partition the module into the state with all the nodes and the moduledefstate,moduledef=model.partition()# verify that the state contains only params, else raise an errorparams,moduledef=model.partition(nnx.Param)# split the state into params and batch_stats, verify no nodes are left(params,batch_stats),moduledef=model.partition(nnx.Param,nnx.BatchStat)# if there are any nodes left, use the `...` filter to capture them(params,batch_stats,rest),moduledef=model.partition(nnx.Param,nnx.BatchStat,...)# using `...` as the only filter is equivalent to not passing any filtersmodel.partition(...)=model.partition()partitionwill make sure all nodes are match by atleast one filter, else it will raise an error. If you have non-Variablenodes likennx.Node,jax.Array, ornumpy.ndarrayattributes, you can use the...filter which will match any node. For a more general filter you can pass a predicate function of the form:(path:Tuple[str,...],value:Any)->boolTo reconstruct the module from a set of substates, you can usemergeas usual but passing the substates as additional arguments:model=moduledef.merge(params,batch_stats,rest)The same is true forapply.y,(state,moduledef)=moduledef.apply(params,batch_stats,rest)(x)Note thatapplywill return a singlestateobject, if you need to re-partition the state you can useState's ownpartitionmethod:params,batch_stats,rest=state.partition(nnx.Param,nnx.BatchStat,...)Alternatively, if you are just interested in a subset of partitions, you can use theState.filtermethod which will not raise an error if some nodes are not matched by any filter:# only get paramsparams=state.filter(nnx.Param)# get params and batch_statsparams,batch_stats=state.filter(nnx.Param,nnx.BatchStat)FiltersFilters let you select subsets of nodes based on some criteria. These are use throughout the API in method likepartition,filter, andpop_state. There are 4 types of filters:type: matches all node instances of the given type....: matches all nodes.(path, any) -> bool: a predicate function that takes a node path and value and returns a boolean.Tuple[Filter, ...]: a tuple of filters, matches all nodes that match any of the filters.NNX also provides the following custom filters:nnx.Not(filter): matches all nodes that do not match the given filternnx.buffers: matches allnumpy.ndarrayandjax.ArraynodesHere is an example of how to useNotandbuffers:rest=module.filter(nnx.Not(nnx.Param))buffers=module.filter(nnx.buffers)Capturing Intermediate ValuesIn NNX you can easily propagate intemediate values by simply assigning them to an attribute at runtime. For convenience, you should assign them to aVariableattribute with acollectionname by usingnnx.varso you can easily retrieve them later.Here is an example of how to create aLinearmodule that captures its output into aVariableattribute with theintermediatescollection name:classLinear(nnx.Module):def__init__(self,din:int,dout:int,*,ctx:nnx.Context):key=ctx.make_rng("params")self.w=nnx.Param(jax.random.uniform(key,(din,dout)))self.b=nnx.Param(jnp.zeros((dout,)))def__call__(self,x):y=x@self.w+self.bself.y=nnx.Intermediate(y)returnymodel=Linear(12,2,ctx=nnx.context(0))Sinceyis only created when the module is called, it is not available upon initialization. However, once you call the moduleywill be created. It is recommended that you usepop_stateto retrieve temporary collections likeIntermediate:y=model(jnp.ones((8,12)))intermediates=model.pop_state(nnx.Intermediate)pop_statewill return aStateobject with the nodes that match the given filter and remove them from the module's attributes.State({
  ('y',): Intermediate(
    value=Array(...)
  )
})If you use the functional API to call the module instead, theIntermediatenodes will be present in the outputstate. To retrieve theintermediatesnodes and optionally separate them from the outputstateyou can useState.partition:state,moduledef=model.partition()y,(state,moduledef)=moduledef.apply(state)(jnp.ones((8,12)))# "pop" the intermediates from the stateintermediates,state=state.partition("intermediates",...)Alternatively, you can useState.filterto retrieve theintermediatesnodes without removing them from thestate.Lifted TransformsNNX lifted transforms analogous versions of JAX transforms but they know how to work with Modules. They usually perform the following tasks:Handle the Module's substates and Context's RNG streams according to the transform's semantics.Properly propagating state in and out of the transform, including updating the input Module's state with updates that happen inside the transform.Here's a diagram illustrating how lifted transformations work:Currently NNX provides thejit,grad, andscanlifted transforms.Manual LiftingIn case you want to use JAX transforms directly you can always use the functional API
to manually lift your Modules.Here we will create an example of how to implement an MLP that uses "scan over layers" to efficiently process a sequence of inputs assuming that each layer has the same parameters and input/output dimensions. The first thing we need to do is create aBlockmodule that represents a single layer, this block with just contain aLinearlayer, aDropoutlayer, and aGELUactivation function:classBlock(nnx.Module):def__init__(self,dim:int,*,ctx:nnx.Context):self.linear=nnx.Linear(dim,dim,ctx=ctx)self.dropout=nnx.Dropout(0.5)def__call__(self,x:jax.Array,*,train:bool,ctx:nnx.Context)->jax.Array:x=self.linear(x)x=self.dropout(x,deterministic=nottrain,ctx=ctx)x=jax.nn.gelu(x)returnxNow we will defineScanMLP. During__init__, instead of creating a list ofBlocks, we will usejax.vmapto create a singleBlockwhose parameters have an addtionallayeraxis. This will allow us to pass the parameters as inputs to scan so it will apply a layer at each step.classScanMLP(nnx.Module):def__init__(self,dim:int,*,n_layers:int,ctx:nnx.Context):params_key=jax.random.split(ctx.make_rng("params"),n_layers)self.n_layers=n_layersself.layers=jax.vmap(lambdakey:Block(dim,ctx=nnx.context(params=key)).partition())(params_key).merge()Note that we split theparamskey inton_layerskeys so each layer has different parameters.Now we will define__call__. Here we need to split thedropoutkey inton_layerskeys so each layer has a different dropout mask, andpartitionthe layers to get theirparams. Bothparamsanddropout_keywill be passed as inputs,xwill be the carry value. Inside thescan_fnwe will merge theparamsback into aBlockmodule and
apply it to the inputx, passing the sliceddropout_keyas part of theContext.def__call__(self,x:jax.Array,*,train:bool,ctx:nnx.Context)->jax.Array:dropout_key=jax.random.split(ctx.make_rng("dropout"),self.n_layers)params,moduledef=self.layers.partition(nnx.Param)defscan_fn(x:inputs):params,dropout_key=inputsmodule=moduledef.merge(params)x=module(x,train=train,ctx=nnx.context(dropout=dropout_key))returnx,module.filter(nnx.Param)x,params=jax.lax.scan(scan_fn,x,(params,dropout_key))self.layers.update_state(params)returnxFinally we applyjax.lax.scan, update thelayersstate with the newparams, and return the finalxvalue.Here is a simple way to test ourScanMLP:model=ScanMLP(10,n_layers=5,ctx=nnx.context(0))x=jnp.ones((3,10))y=model(x,train=True,ctx=nnx.context(dropout=1))For a more robust implementation with comments take a look at theScan over layersexample.Case StudiesShared StateIn NNX, you can create modules that share state between them. This is useful when designing complex neural network architectures, as it allows you to reuse certain layers and reduce the number of learnable parameters.Here's an example of creating a module with shared state:classBlock(nnx.Module):def__init__(self,linear:nnx.Linear,*,ctx:nnx.Context):self.linear=linearself.bn=nnx.BatchNorm(2,ctx=ctx)def__call__(self,x,*,ctx:nnx.Context):x=self.linear(x)x=self.bn(x,ctx=ctx)x=nnx.relu(x)returnxclassModel(nnx.Module):def__init__(self,*,ctx:nnx.Context):shared=nnx.Linear(2,2,ctx=ctx)self.block1=Block(shared,ctx=ctx)self.block2=Block(shared,ctx=ctx)def__call__(self,x,*,ctx:nnx.Context):x=self.block1(x,ctx=ctx)x=self.block2(x,ctx=ctx)returnxIn this example, theModelmodule contains two instances of theBlockmodule. Each instance shares the samennx.Linearmodule. To run the model, you can use the Contextflagsargument to set theuse_running_averageflag for allBatchNormmodules.Here's an example of computing the loss for aModelinstance:defloss_fn(model:Model,x:jax.Array,y:jax.Array):ctx=nnx.context(flags=dict(use_running_average=True))y_pred=model(x,ctx=ctx)returnjnp.mean((y-y_pred)**2)It's important to note that the state for the sharednnx.Linearmodule will be kept in sync at all times on bothBlockinstances, including during gradient updates.