JAX

What it is

Google's high-performance ML library — NumPy with autograd, JIT compilation (XLA), and easy parallelism (pmap, vmap). The framework behind Gemini and most modern Google research.

How Vaaani uses it

• Research-scale training on TPU or GPU clusters
• Custom gradient computations for non-standard objectives
• Vectorized scientific computing (physics, biology models)
• Functional purity makes large-scale training reproducible

Why it makes the cut

When throughput on a multi-GPU box matters and the team is comfortable with functional programming, JAX is unbeaten.

Sample code

import jax, jax.numpy as jnp
from jax import grad, jit

def loss(w, x, y):
    return jnp.mean((jnp.dot(x, w) - y) ** 2)

grad_loss = jit(grad(loss))