Update environments and dependencies.

PiperOrigin-RevId: 632263761

experiments/environment.yml

@@ -0,0 +1,25 @@
+name: tfm_env
+
+channels:
+  - conda-forge
+  - defaults
+  - anaconda
+dependencies:
+  - jupyterlab
+  - pip
+  - python=3.10
+  - pip:
+    - datasetsforecast
+    - fire
+    - git+https://github.com/awslabs/gluon-ts.git
+    - huggingface_hub[cli]
+    - neuralforecast
+    - orjson
+    - statsforecast
+    - utilsforecast
+    - 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]"
+    - einshape

experiments/environment_cpu.yml

@@ -0,0 +1,24 @@
+name: tfm_env
+
+channels:
+  - conda-forge
+  - defaults
+  - anaconda
+dependencies:
+  - jupyterlab
+  - pip
+  - python=3.10
+  - pip:
+    - datasetsforecast
+    - fire
+    - git+https://github.com/awslabs/gluon-ts.git
+    - huggingface_hub[cli]
+    - neuralforecast
+    - orjson
+    - statsforecast
+    - utilsforecast
+    - git+https://github.com/amazon-science/chronos-forecasting.git
+    - praxis
+    - paxml
+    - jax[cpu]
+    - einshape