Update environments and dependencies.

PiperOrigin-RevId: 632263761

README.md

@@ -30,9 +30,9 @@ Please look into the README files in the respective benchmark directories within
 
 ## Installation
 
-We have two environment files. For GPU installation (assuming CUDA 12 has been
-setup), you can create a conda environment `tfm_env` from the base folder
-through:
+For calling TimesFM, We have two environment files. Inside `timesfm`, for
+GPU installation (assuming CUDA 12 has been setup), you can create a conda
+environment `tfm_env` from the base folder through:
 
 ```
 conda env create --file=environment.yml
@@ -53,6 +53,14 @@ pip install -e .
 ```
 to install the package.
 
+**Note**: 
+
+1. Running the provided benchmarks would require additional dependencies.
+Please use the environment files under `experiments` instead.
+
+2. The dependency `lingvo` does not support `pip install` for macOS. For now you
+can manually install it from the source following [these instructions](https://github.com/tensorflow/lingvo).
+
 ## Usage 
 
 ### Initialize the model and load a checkpoint.
@@ -112,7 +120,7 @@ Array inputs, with the frequencies set to low, medium and high respectively.
 ```python
 import numpy as np
 forecast_input = [
-    np.sin(np.linspace(0, 20, 100))
+    np.sin(np.linspace(0, 20, 100)),
     np.sin(np.linspace(0, 20, 200)),
     np.sin(np.linspace(0, 20, 400)),
 ]

environment.yml

@@ -9,15 +9,7 @@ dependencies:
   - 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  

environment_cpu.yml

@@ -9,15 +9,7 @@ dependencies:
   - 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]

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

pyproject.toml

@@ -5,6 +5,7 @@ name = "timesfm"
 description = "Open weights time-series foundation model from Google Research."
 version = "0.0.1"
 dependencies = [
+    "einshape>=1.0.0",
     "paxml>=1.4.0",
     "praxis>=1.4.0",
     "jax>=0.4.26",