From a4beacd02318fab3e4b47fd7ebfc6e02bad95294 Mon Sep 17 00:00:00 2001 From: Rajat Sen Date: Tue, 21 May 2024 20:02:57 +0000 Subject: [PATCH] Pinning jax version. --- environment.yml | 5 ++--- environment_cpu.yml | 2 +- experiments/environment.yml | 3 +-- experiments/environment_cpu.yml | 2 +- pyproject.toml | 2 +- 5 files changed, 6 insertions(+), 8 deletions(-) diff --git a/environment.yml b/environment.yml index a381a0d..dd5388b 100644 --- a/environment.yml +++ b/environment.yml @@ -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 diff --git a/environment_cpu.yml b/environment_cpu.yml index 61f1ebe..772c4b7 100644 --- a/environment_cpu.yml +++ b/environment_cpu.yml @@ -13,5 +13,5 @@ dependencies: - utilsforecast - praxis - paxml - - jax[cpu] + - jax[cpu]==0.4.26 - einshape diff --git a/experiments/environment.yml b/experiments/environment.yml index 9238669..25b4b7b 100644 --- a/experiments/environment.yml +++ b/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 diff --git a/experiments/environment_cpu.yml b/experiments/environment_cpu.yml index dabd911..2301b8e 100644 --- a/experiments/environment_cpu.yml +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 2a34f2e..9c097bb 100644 --- a/pyproject.toml +++ b/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", ]