JAX
Also known as JAX framework, jax.jit, Google JAX
In plain English
JAX is Google’s Python framework for array computing that compiles functions through XLA and has been the native way to program TPUs.
Technical definition
JAX is a Python library for composable function transformations such as jit, grad, and vmap over NumPy-style arrays, compiled through XLA and used as the primary first-party framework for TPUs.
Engineering details
JAX programs are functional: state such as model weights and the KV cache is passed in and out explicitly so jax.jit can trace a step into a graph. That model fits XLA well and made JAX the mature path for TPU primitives and parallelism, which is why the first external vLLM TPU backend translated PyTorch into JAX through TorchAX and why SGLang-JAX exists as a separate JAX-native engine. TorchTPU changes the relationship: PyTorch becomes the user-facing framework, but JAX-backed custom kernels and Pallas remain callable underneath, and TPU-Sync works natively with both JAX and TorchTPU.
Why it matters
The open-source serving ecosystem is written in PyTorch. Requiring engines to cross a PyTorch-to-JAX boundary slowed feature parity and model bring-up on TPU, so Google is moving the serving stack to PyTorch-native while keeping JAX and Pallas for kernels and for its own internal workloads.
How to read it in InferenceX
Initial Qwen3.5 TPU support added pure JAX implementations of causal Conv1D and Gated DeltaNet before later Pallas kernels optimized them. The InferenceX Official Preview numbers come from the TorchTPU vLLM stack rather than the JAX-translated tpu-inference backend.
Source material
See the concept in real benchmarks
TPU Inference Externalization Full Steam Ahead
InferenceX, Up to 50% Better Performance per Dollar, Rapid Externalization of TPU stack, Growing Customer Base, Ironwood, TPUv8i, Reducing CUDA Moat
InferenceX v2: NVIDIA Blackwell Vs AMD vs Hopper - Formerly InferenceMAX
GB300 NVL72, MI355X, B200, H100, Disaggregated Serving, Wide Expert Parallelism, Large Mixture of Experts, SGLang, vLLM, TRTLLM