JAX
Also: jit, vmap, pmap
A numerical library built around transforming pure functions rather than objects that hold state: jit compiles one, grad differentiates it, vmap batches it, and pmap or shard_map spreads it across devices. Used heavily at some frontier labs, and the native stack for TPUs.
What people get wrong
Not a replacement for PyTorch in general. Most research code, most tutorials and most production serving are PyTorch, and learning JAX first will slow you down unless the job you want is written in it.
A definition is not an understanding. This is taught properly in PyTorch in practice, part of PyTorch.
Learn this properlyTerms from the same topic