J jax.readthedocs.io JAX – NumPy on the CPU, GPU, and TPU 276 points peter_d_sherman 3 years ago 143 comments