jax 0.3.9-foss-2021a-CUDA-11.3.1

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Accessing jax 0.3.9-foss-2021a-CUDA-11.3.1

To load the module for jax 0.3.9-foss-2021a-CUDA-11.3.1 please use this command on Baskerville:

module load bask-apps/live
module load jax/0.3.9-foss-2021a-CUDA-11.3.1

Bask Apps Env

live

Extensions

  • jax 0.3.9
  • opt_einsum 3.3.0

More Information

For more information visit the jax website.

Dependencies

This version of jax has a direct dependency on: CUDA/11.3.1 cuDNN/8.2.1.32-CUDA-11.3.1 flatbuffers-python/2.0-GCCcore-10.3.0 foss/2021a NCCL/2.10.3-GCCcore-10.3.0-CUDA-11.3.1 Python/3.9.5-GCCcore-10.3.0 SciPy-bundle/2021.05-foss-2021a typing-extensions/3.10.0.0-GCCcore-10.3.0

Required By

This version of jax is a direct dependent of: AlphaFold/2.2.2-foss-2021a-CUDA-11.3.1

Last modified on 22nd August 2022