JAX is Autograd and XLA, brought together for high-performance machine learning research.

More Information

For more information visit the JAX website.

Available Versions

These versions of JAX are available on the BEAR systems (BlueBEAR, BEARCloud VMs, and CaStLeS VMs). These will be retained in accordance with our Applications Support and Retention Policy.

Version BEAR Apps Version
0.1.77-fosscuda-2019b-Python-3.7.4 2019b