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:
- conda-forge
@@ -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
+1 -1
View File
@@ -13,5 +13,5 @@ dependencies:
- utilsforecast
- praxis
- paxml
- jax[cpu]
- jax[cpu]==0.4.26
- einshape
+1 -2
View File
@@ -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
+1 -1
View File
@@ -20,5 +20,5 @@ dependencies:
- git+https://github.com/amazon-science/chronos-forecasting.git
- praxis
- paxml
- jax[cpu]
- jax[cpu]==0.4.26
- einshape
+1 -1
View File
@@ -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",
]