Optax

[lexming]

• link to support ticket: INC0149937 (VUB)
• website: https://github.com/google-deepmind/optax
• installation docs: https://optax.readthedocs.io
• toolchain: foss/2023a
• easyblock to use: PythonBundle
• required dependencies:
  • absl-py>=0.7.1
  • chex>=0.1.8
  • jax>=0.4.13
  • jaxlib>=0.4.13
  • numpy>=1.25.0
• effort: seems easy
• other install methods
  • conda: no
  • container image: no
  • pre-built binaries (RHEL8 Linux x86_64): no
  • easyconfig outside EasyBuild: no

[lexming]

PRs:

• reorganize Flax/JAX stack in 2023a: move jax + Optax to gfbf/2023a toolchain + use standalone Flax + absl-py as dependencies easybuilders/easybuild-easyconfigs#21038
• use stand-alone absl-py as dependency for jax w/ gfbf/2023a + CUDA 12.1.1 + add Optax v0.2.2, Flax v0.8.4, scvi-tools v1.1.2 easybuilders/easybuild-easyconfigs#21039