Package: pax3
Introduction|Getting started|Functional programming|Examples|Modules|Fine-tuningIntroductionPAX is aJAX-based library for training neural networks.PAX modules are registered as JAXpytree, therefore, they can be input or output of JAX transformations such asjax.jit,jax.grad, etc. This makes programming with modules very convenient and easy to understand.InstallationInstall from PyPI:pipinstallpax3Or install the latest version from Github:pipinstallgit+https://github.com/ntt123/pax.git## or test mode to run tests and examplespipinstallgit+https://github.com/ntt123/pax.git#egg=pax3[test]Getting startedBelow is a simple example of aLinearmodule.importjax.numpyasjnpimportpaxclassLinear(pax.Module):weight:jnp.ndarraybias:jnp.ndarrayparameters=pax.parameters_method("weight","bias")def__init__(self):super().__init__()self.weight=jnp.array(0.0)self.bias=jnp.array(0.0)def__call__(self,x):returnself.weight*x+self.biasThe implementation is very similar to a normal python class. However, we need an additional lineparameters=pax.parameters_method("weight","bias")to declare thatweightandbiasaretrainable parametersof the Linear module.PAX functional programmingpax.pureA PAX module can have internal states. For example, below is a simpleCountermodule with an internal counter.classCounter(pax.Module):count:jnp.ndarraydef__init__(self):super().__init__()self.count=jnp.array(0)def__call__(self):self.count=self.count+1returnself.countHowever, PAXaimsto guarantee that modules will have no side effects from the outside point of view.
Therefore, the modifications of these internal states are restricted. For example, we get an error when trying to callCounterdirectly.counter=Counter()count=counter()# ...# ----> 9         self.count = self.count + 1# ...# ValueError: Cannot modify a module in immutable mode.# Please do this computation inside a function decorated by `pax.pure`.Only functions decorated bypax.pureare allowed to modify input module's internal states.@pax.puredefupdate_counter(counter:Counter):count=counter()returncounter,countcounter,count=update_counter(counter)print(counter.count,count)# 1 1Note that we have to returncounterin the output ofupdate_counter, otherwise, thecounterobject will not be updated. This is becausepax.pureonly providesupdate_countera copy of thecounterobject.pax.purecallFor convenience, PAX provides thepax.purecallfunction.
It is a shortcut forpax.pure(lambda f, x: [f, f(x)]).Instead of implementing anupdate_counterfunction, we can do the same thing with:counter,count=pax.purecall(counter)print(counter.count,count)# 2, 2Replacing partsPAX provides utility methods to modify a module in a functional way.Thereplacemethod creates a new module with attributes replaced.
For example, to replaceweightandbiasof apax.Linearmodule:fc=pax.Linear(2,2)fc=fc.replace(weight=jnp.ones((2,2)),bias=jnp.zeros((2,)))Thereplace_nodemethod replaces a pytree node of a module:f=pax.Sequential(pax.Linear(2,3),pax.Linear(3,4),)f=f.replace_node(f[-1],pax.Linear(3,5))print(f.summary())# Sequential# ├── Linear(in_dim=2, out_dim=3, with_bias=True)# └── Linear(in_dim=3, out_dim=5, with_bias=True)PAX and other librariesPAX learns a lot from other libraries:PAX borrows the idea thata module is also a pytreefromtreexandequinox.PAX uses the concept oftrainable parametersandnon-trainable statesfromdm-haiku.PAX has similar methods to PyTorch such asmodel.apply(),model.parameters(),model.eval(), etc.PAX usesobjax's approach to implement optimizers as modules.PAX usesjmplibrary for supporting mixed precision.And of course, PAX is heavily influenced byjaxfunctional programming approach.ExamplesA good way to learn aboutPAXis to see examples in theexamples/directory.Click to expandPathDescriptionchar_rnn.pytrain a RNN language model on TPU.transformer/train a Transformer language model on TPU.mnist.pytrain an image classifier onMNISTdataset.notebooks/VAE.ipynbtrain a variational autoencoder.notebooks/DCGAN.ipynbtrain a DCGAN model onCeleb-Adataset.notebooks/fine_tuning_resnet18.ipynbfinetune a pretrained ResNet18 model oncats vs dogsdataset.notebooks/mixed_precision.ipynbtrain a U-Net image segmentation with mixed precision.mnist_mixed_precision.pytrain an image classifier with mixed precision.wave_gru/train a WaveGRU vocoder: convert mel-spectrogram to waveform.denoising_diffusion/train a denoising diffusion model onCeleb-Adataset.ModulesAt the moment, PAX includes:pax.Embed,pax.Linear,pax.{GRU, LSTM},pax.{BatchNorm1D, BatchNorm2D, LayerNorm, GroupNorm},pax.{Conv1D, Conv2D, Conv1DTranspose, Conv2DTranspose},pax.{Dropout, Sequential, Identity, Lambda, RngSeq, EMA}.OptimizersPAX has its optimizers implemented in a separate libraryopax. Theopaxlibrary supports many common optimizers such asadam,adamw,sgd,rmsprop. Visit opax's GitHub repository for more information.Fine-tunning modelsPAX's Module provides thepax.freeze_parameterstransformation to convert all trainable parameters to non-trainable states.net=pax.Sequential(pax.Linear(28*28,64),jax.nn.relu,pax.Linear(64,10),)net=pax.freeze_parameters(net)net=net.set(-1,pax.Linear(64,2))After this,net.parameters()will only return trainable parameters of the last layer.