Pinning jax version.

This commit is contained in:
Rajat Sen
2024-05-21 20:02:57 +00:00
parent 469b7071dc
commit a4beacd023
5 changed files with 6 additions and 8 deletions
+2 -3
View File
@@ -1,4 +1,4 @@
name: tfm_env name: tfm_env_t2
channels: channels:
- conda-forge - conda-forge
@@ -13,6 +13,5 @@ dependencies:
- utilsforecast - utilsforecast
- praxis - praxis
- paxml - paxml
- -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html - jax[cuda12]==0.4.26
- "jax[cuda12_pip]"
- einshape - einshape
+1 -1
View File
@@ -13,5 +13,5 @@ dependencies:
- utilsforecast - utilsforecast
- praxis - praxis
- paxml - paxml
- jax[cpu] - jax[cpu]==0.4.26
- einshape - einshape
+1 -2
View File
@@ -20,6 +20,5 @@ dependencies:
- git+https://github.com/amazon-science/chronos-forecasting.git - git+https://github.com/amazon-science/chronos-forecasting.git
- praxis - praxis
- paxml - paxml
- -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html - jax[cuda12]==0.4.26
- "jax[cuda12_pip]"
- einshape - einshape
+1 -1
View File
@@ -20,5 +20,5 @@ dependencies:
- git+https://github.com/amazon-science/chronos-forecasting.git - git+https://github.com/amazon-science/chronos-forecasting.git
- praxis - praxis
- paxml - paxml
- jax[cpu] - jax[cpu]==0.4.26
- einshape - einshape
+1 -1
View File
@@ -8,7 +8,7 @@ dependencies = [
"einshape>=1.0.0", "einshape>=1.0.0",
"paxml>=1.4.0", "paxml>=1.4.0",
"praxis>=1.4.0", "praxis>=1.4.0",
"jax>=0.4.26", "jax==0.4.26",
"numpy>=1.26.4", "numpy>=1.26.4",
"pandas>=2.1.4", "pandas>=2.1.4",
] ]