Updating compile file with required JAX version to fix the amazon p…

…lugin
PennyLaneAI · Jan 31, 2025 · b7e36dc · b7e36dc
1 parent 8fa68f4
commit b7e36dc
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/.github/workflows/braket-latest-latest.yml b/.github/workflows/braket-latest-latest.yml
diff --git a/.github/workflows/braket-latest-rc.yml b/.github/workflows/braket-latest-rc.yml
diff --git a/compile.py b/compile.py
@@ -73,9 +73,14 @@
         "additional_setup": dedent("""
             - name: Install TF
               run: |
-                pip install tensorflow~=$TF_VERSION keras~=$TF_VERSION"""
+                pip install tensorflow~=$TF_VERSION keras~=$TF_VERSION
+                                   
+            - name: Install JAX
+              run: |
+                pip install jax==$JAX_VERSION jaxlib==$JAX_VERSION"""
+
         ),
-        "additional_env_vars": "TF_VERSION: 2.12.0\n  TORCH_VERSION: 2.0.0+cpu",
+        "additional_env_vars": "TF_VERSION: 2.12.0\n  TORCH_VERSION: 2.0.0+cpu\n  JAX_VERSION: 0.4.28",
         "no_deprecation_error": True,
     },
 ]