Merge pull request #265 from Killer3048/feature/lazy-import-xreg

Implement Lazy import for `xreg` dependencies in `TimesFM`

README.md

@@ -78,18 +78,23 @@ poetry install -E  torch
 
 After than you can run the timesfm under `poetry shell` or do `poetry run python3 ...`.
 
-**Note**: 
+**Additional Note**: 
 
-1. Running the provided benchmarks would require additional dependencies.
+If you plan to use the **`forecast_with_covariates`** function (which requires external regressors), 
-Please see the `experiments` section fro more instructions.
+you need to install **JAX** and **jaxlib**. If you installed the base version of TimesFM (`torch`), you must manually install the dependencies for **`forecast_with_covariates`** support:
+```
+pip install jax jaxlib
+```
 
-2. The dependency `lingvo` does not support ARM architectures, and the code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned.
+**Why is this needed?**  
+The `forecast_with_covariates` method relies on the `xreg_lib` module, which depends on JAX and jaxlib. If these packages are not installed, 
+calling `forecast_with_covariates` will raise an error. However, due to a lazy import mechanism, `xreg_lib` (and hence JAX/jaxlib) is not needed for standard `forecast` calls.
 
 ### Notes
 
 1. Running the provided benchmarks would require additional dependencies. Please see the `experiments` folder.
 
-2. The dependency `lingvo` does not support ARM architectures, and the PAX version is not working for machines with Apple silicon.
+2. The dependency `lingvo` does not support ARM architectures, and the code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned.
 
 ### Install from PyPI (and publish)
 
@@ -241,6 +246,12 @@ forecast_df = tfm.forecast_on_df(
 
 We now have an external regressors library on top of TimesFM that can support static covariates as well as dynamic covariates available in the future. We have an usage example in [notebooks/covariates.ipynb](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb).
 
+If you plan to use the **`forecast_with_covariates`** on timesfm `torch` version, you need to install **JAX** and **jaxlib**. 
+You must manually install the dependencies for **`forecast_with_covariates`** support:
+```
+pip install jax jaxlib
+```
+
 Let's take a toy example of forecasting sales for a grocery store: 
 
 **Task:** Given the observed the daily sales of this week (7 days), forecast the daily sales of next week (7 days).

pyproject.toml

@@ -42,7 +42,7 @@ absl-py = ">=1.4.0"
 
 [tool.poetry.extras]
 pax = ["paxml", "lingvo", "jax", "jaxlib"]
-torch = ["torch",  "jax", "jaxlib"]  # jax & jaxlib are already in pax
+torch = ["torch"] 
 
 [tool.poetry.dependencies.paxml]
 version = ">=1.4.0"

src/timesfm/timesfm_base.py

@@ -17,17 +17,20 @@ import collections
 import dataclasses
 import logging
 import multiprocessing
-from typing import Any, Literal, Sequence
+from typing import Any, Literal, Sequence, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
 
 from utilsforecast.processing import make_future_dataframe
 
+if TYPE_CHECKING:
     from . import xreg_lib
-
     Category = xreg_lib.Category
     XRegMode = xreg_lib.XRegMode
+else:
+    Category = int | str
+    XRegMode = str
 
 _TOL = 1e-6
 DEFAULT_QUANTILES = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
@@ -464,6 +467,8 @@ class TimesFmBase:
       the outputs of the xreg.
     """
 
+    from . import xreg_lib
+
     # Verify and bookkeep covariates.
     if not (dynamic_numerical_covariates or dynamic_categorical_covariates or
             static_numerical_covariates or static_categorical_covariates):