Merge pull request #45 from google-research/rajat_dev

Pining jax version.

environment.yml

@@ -13,6 +13,5 @@ dependencies:
     - utilsforecast
     - praxis
     - paxml
-    - -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  
-    - "jax[cuda12_pip]"
+    - jax[cuda12]==0.4.26
     - einshape

environment_cpu.yml

@@ -13,5 +13,5 @@ dependencies:
     - utilsforecast
     - praxis
     - paxml
-    - jax[cpu]
+    - jax[cpu]==0.4.26
     - einshape

experiments/environment.yml

@@ -20,6 +20,5 @@ dependencies:
     - git+https://github.com/amazon-science/chronos-forecasting.git
     - praxis
     - paxml
-    - -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html  
-    - "jax[cuda12_pip]"
+    - jax[cuda12]==0.4.26
     - einshape

experiments/environment_cpu.yml

@@ -20,5 +20,5 @@ dependencies:
     - git+https://github.com/amazon-science/chronos-forecasting.git
     - praxis
     - paxml
-    - jax[cpu]
+    - jax[cpu]==0.4.26
     - einshape

pyproject.toml

@@ -8,7 +8,7 @@ dependencies = [
     "einshape>=1.0.0",
     "paxml>=1.4.0",
     "praxis>=1.4.0",
-    "jax>=0.4.26",
+    "jax==0.4.26",
     "numpy>=1.26.4",
     "pandas>=2.1.4",
 ]