JAX

JAX is an open-source numerical computing library developed by Google and first released in 2018. It brings together a familiar NumPy-style programming interface with three powerful transformations: automatic differentiation, just-in-time (JIT) compilation, and automatic vectorisation. JAX has become a key framework for training state-of-the-art foundation models, used by organisations including Google DeepMind, Anthropic, xAI, and Apple.

JAX is built on a functional programming model. Rather than mutating state in place, programs are written as pure mathematical functions that JAX can analyse and transform. This design is what allows the library to apply its transformations composably and reliably.

Core transformations

The function grad performs automatic differentiation, returning a new function that computes the gradient of the original. Because it composes, higher-order derivatives are obtained by applying it repeatedly. The function jit compiles Python functions into optimised machine code using the XLA (Accelerated Linear Algebra) compiler, often yielding large speed-ups by fusing operations and reducing overhead. The function vmap automatically vectorises a function so that it operates over batches without the developer writing explicit loops, and pmap and related tools distribute computation across multiple devices.

These transformations run unchanged on central processing units, GPUs, and Google's Tensor Processing Units, making JAX especially well suited to large-scale training on TPU pods.

The JAX AI stack

JAX itself is deliberately minimal, providing the numerical core while higher-level libraries supply the conveniences needed for full machine-learning workflows. Google has packaged these into the JAX AI Stack, an end-to-end platform co-designed with Cloud TPUs. Flax provides a flexible API for authoring neural-network models. Optax offers composable gradient-processing and optimisation transformations. Orbax handles asynchronous, distributed checkpointing so that long training runs survive hardware failures. Newer additions announced in 2025 include Metrax for efficient evaluation metrics and JAX-Privacy 1.0 for differentially private training pipelines.

Position relative to PyTorch and TensorFlow

JAX occupies a distinct niche. Where PyTorch emphasises an imperative, easy-to-debug style and TensorFlow emphasises production deployment, JAX targets high-performance research and large-scale training where its functional purity and XLA compilation deliver strong performance. The trade-off is a steeper learning curve and a smaller, though rapidly growing, ecosystem. Its functional approach also extends beyond deep learning into scientific computing, physics simulation, and other fields that benefit from composable differentiation.

| Transformation | Purpose | |----------------|---------| | grad | Automatic differentiation | | jit | Just-in-time compilation via XLA | | vmap | Automatic batching / vectorisation | | pmap | Parallelism across devices |

JAX is most relevant to the advanced end of Malaysia's AI activity: research groups, well-resourced startups, and any effort to train or fine-tune large foundation models. Its tight integration with Tensor Processing Units matters as Malaysia expands its computing infrastructure, with major data-centre investments in Johor and the Klang Valley by hyperscale cloud providers offering access to TPU and GPU capacity in the region.

For Malaysian initiatives exploring sovereign AI and locally developed large language models, such as efforts to build Malay-language and multilingual models suited to the national context, JAX provides the performance characteristics needed for efficient large-scale training. The framework's adoption by leading model developers makes familiarity with it valuable for researchers at institutions such as Universiti Malaya, Universiti Sains Malaysia, and MIMOS, the national applied-research agency.

The JAX-Privacy library is notable for Malaysian organisations subject to the Personal Data Protection Act (PDPA), as it enables training models with formal privacy guarantees, a capability of interest to banking, healthcare, and government applications handling sensitive data.

JAX's relatively specialised nature means it is more likely to appear in postgraduate research and frontier industry projects than in general SME adoption, where PyTorch and scikit-learn remain more common. As Malaysia's AI talent pool and compute capacity grow under the MyDigital Blueprint and the National AI agenda, demand for high-performance frameworks like JAX is expected to rise.

1. Bradbury, J., et al. (2018). JAX: Composable Transformations of Python and NumPy Programs. Google.
2. Google Developers Blog. (2025). Building Production AI on Google Cloud TPUs with JAX. developers.googleblog.com.
3. Google Research. (2025). Differentially Private Machine Learning at Scale with JAX-Privacy. research.google.
4. Google Cloud Blog. (2025). A Guide to JAX for PyTorch Developers. cloud.google.com.