Package: jaxnet
JAXnetJAXnet is a deep learning library based onJAX.
JAXnet's functional API provides unique benefits over TensorFlow2, Keras and PyTorch,
while maintaining user-friendliness, modularity and scalability:More robustness through immutable weights, no global compute graph.GPU-compilednumpycode for networks, training loops, pre- and postprocessing.Regularization and reparametrization of any module or whole networks in one line.No global random state, flexible random key control.If you already know stax, readthis.Modularitynet=Sequential(Dense(1024),relu,Dense(1024),relu,Dense(4),log_softmax)creates a neural net model from predefined modules.ExtensibilityDefine your own modules using@parametrizedfunctions. You can reuse other modules:fromjaximportnumpyasjnp@parametrizeddefloss(inputs,targets):return-jnp.mean(net(inputs)*targets)All modules are composed in this way.jax.numpyis mirroringnumpy,
meaning that if you know how to usenumpy, you know most of JAXnet.
Compare this to TensorFlow2/Keras:importtensorflowastffromtensorflow.kerasimportSequentialfromtensorflow.keras.layersimportDense,Lambdanet=Sequential([Dense(1024,'relu'),Dense(1024,'relu'),Dense(4),Lambda(tf.nn.log_softmax)])defloss(inputs,targets):return-tf.reduce_mean(net(inputs)*targets)Notice howLambdalayers are not needed in JAXnet.reluandlogsoftmaxare plain Python functions.Immutable weightsDifferent from TensorFlow2/Keras, JAXnet has no global compute graph.
Modules likenetandlossdo not contain mutable weights.
Instead, weights are contained in separate, immutable objects.
They are initialized withinit_parameters, provided example inputs and a random key:fromjax.randomimportPRNGKeydefnext_batch():returnjnp.zeros((3,784)),jnp.zeros((3,4))params=loss.init_parameters(*next_batch(),key=PRNGKey(0))print(params.sequential.dense2.bias)# [-0.01101029, -0.00749435, -0.00952365,  0.00493979]Instead of mutating weights inline, optimizers return updated versions of weights.
They are returned as part of a new optimizer state, and can be retrieved viaget_parameters:opt=optimizers.Adam()state=opt.init(params)for_inrange(10):state=opt.update(loss.apply,state,*next_batch())# accelerate with jit=Truetrained_params=opt.get_parameters(state)applyevaluates a network:test_loss=loss.apply(trained_params,*test_batch)# accelerate with jit=TrueGPU support and compilationJAX allows functionalnumpy/scipycode to be accelerated.
Make it run on GPU by replacing yournumpyimport withjax.numpy.
Compile a function by decorating it withjit.
This will free your function from slow Python interpretation, parallelize operations where possible and optimize your compute graph.
It provides speed and scalability at the level of TensorFlow2 or PyTorch.Due to immutable weights, whole training loops can be compiled / run on GPU (demo).jitwill make your training as fast as mutating weights inline, and weights will not leave the GPU.
You can write functional code without worrying about performance.You can easily acceleratenumpy/scipypre-/postprocessing code in the same way (demo).Regularization and reparametrizationIn JAXnet, regularizing a model can be done in one line (demo):loss=L2Regularized(loss,scale=.1)lossis now just another module that can be used as above.
Reparametrized layers are one-liners, too (seeAPI).
JAXnet allows regularizing or reparametrizing any module or subnetwork without changing its code.
This is possible because modules do not instantiate any variables.
Instead each module provides a function (apply) with parameters as an argument.
This function can be wrapped to build layers likeL2Regularized.In contrast, TensorFlow2/Keras/PyTorch have mutable variables baked into their model API. They therefore require:Regularization arguments on layer level, with separate code necessary for each layer.Reparametrization arguments on layer level, and separate implementations foreachlayer.Random key controlJAXnet does not have global random state.
Random keys are provided explicitly, making code deterministic and independent of previously executed code by default.
This can help debugging and is more flexible (demo).
Read more on random numbers in JAXhere.Step-by-step debuggingJAXnet allows step-by-step debugging with concrete values like any plain Python function
(whenjitcompilation is not used).API and demosFind more details on the APIhere.See JAXnet in action in your browser:Mnist Classifier,Mnist VAE,OCR with RNNs,ResNet,WaveNet,PixelCNN++andPolicy Gradient RL.InstallationThis is a preview. Expect breaking changes!Python 3.6 or higher is supported. Install withpip3 install jaxnetTo use GPU, first install theright version of jaxlib.QuestionsPlease feel free to create an issue on GitHub.