[tough]

There's a JAX for AI/LM too

https://github.com/jax-ml/jax

but yeah no idea which the OP meant