jax¶
Description¶
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
More information¶
- Homepage: https://pypi.python.org/pypi/jax
Included extensions¶
etils-0.8.0, jax-0.3.23, opt_einsum-3.3.0
Availability¶
| Module | NIC5 | Hercules2 | Dragon2 | Lemaitre3 | Dragon1 |
|---|---|---|---|---|---|
| jax/0.2.19-foss-2020b | ✓ | ✓ | |||
| jax/0.2.19-fosscuda-2020b | ✓ | ✓ | |||
| jax/0.3.23-foss-2021b-CUDA-11.4.1 | ✓ | ✓ | |||
| jax/0.3.25-foss-2022a-CUDA-11.7.0 | ✓ |