py-alphafold: update to 2.2.4, update dependencies (#33876)

* py-alphafold: update to 2.2.4, update dependencies * style
2022-12-02 17:12:06 -05:00
parent c10b84f08d
commit 588d2e295f
10 changed files with 92 additions and 32 deletions
--- a/share/spack/gitlab/cloud_pipelines/stacks/ml-cpu/spack.yaml
+++ b/share/spack/gitlab/cloud_pipelines/stacks/ml-cpu/spack.yaml
@@ -28,9 +28,8 @@ spack:
    - py-transformers

    # JAX
-    # https://github.com/google/jax/issues/12614
-    # - py-jax
-    # - py-jaxlib
+    - py-jax
+    - py-jaxlib

    # Keras
    - py-keras
--- a/share/spack/gitlab/cloud_pipelines/stacks/ml-cuda/spack.yaml
+++ b/share/spack/gitlab/cloud_pipelines/stacks/ml-cuda/spack.yaml
@@ -31,9 +31,8 @@ spack:
    - py-transformers

    # JAX
-    # https://github.com/google/jax/issues/12614
-    # - py-jax
-    # - py-jaxlib
+    - py-jax
+    - py-jaxlib

    # Keras
    - py-keras
--- a/share/spack/gitlab/cloud_pipelines/stacks/ml-rocm/spack.yaml
+++ b/share/spack/gitlab/cloud_pipelines/stacks/ml-rocm/spack.yaml
@@ -33,9 +33,8 @@ spack:
    - py-transformers

    # JAX
-    # https://github.com/google/jax/issues/12614
-    # - py-jax
-    # - py-jaxlib
+    - py-jax
+    - py-jaxlib

    # Keras
    - py-keras