Gradient computation#

Easy gradient calculation is a key feature of JAX.In theJaxOpt library value and grad is directly utilized for users in multiple optimization algorithms inits source code.

Similarly the same Dynamax Optax pairing mentioned above is an example ofgradients enabling estimation methods that were challenging historicallyMaximum Likelihood Expectation using Optax.

Computational speedup on a single core across multiple devices#

Models defined in JAX can then be compiled to enable single computation speedup through JIT compiling.The same compiled code can then be sent to a CPU device,to a GPU or TPU device for additional speedup,typically with no additional changes needed.This allows for a smooth workflow from development into production.In Dynamax the computationally expensive portion of a Linear State Space Model solver has beenjitted.A more complex example comes from PyTensor which compiles a JAX function dynamically and thenjits the constructed function.

Single and multi computer speedup using parallelization#

Another benefit of JAX is the simplicity of parallelizing computation usingpmap andvmap function calls or decorators.In Dynamax state space models are parallelized with aVMAP decoratora practical example of this use case being multi object tracking.

Incorporating JAX code into your, or your users, workflows#

JAX is quite composable and can be used in multiple ways.JAX can be used with a standalone pattern, where the user defines all the calculations themselves.However other patterns, such as using libraries built on jax that provide specific functionality.These can be libraries that define specific types of models,such as Neural Networks or State Space models or others,or provide specific functionality such as optimization.Here are more specific examples of each pattern.