From 9a520bb9406d6a7534954848f23d9c69e763f5ad Mon Sep 17 00:00:00 2001 From: Killer3048 Date: Wed, 9 Apr 2025 12:55:04 +0700 Subject: [PATCH 1/5] Implement lazy import for xreg dependencies in forecast_with_covariates to avoid unnecessary JAX installation --- src/timesfm/timesfm_base.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/timesfm/timesfm_base.py b/src/timesfm/timesfm_base.py index 088cfab..4502d21 100644 --- a/src/timesfm/timesfm_base.py +++ b/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 -from . import xreg_lib - -Category = xreg_lib.Category -XRegMode = xreg_lib.XRegMode +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) @@ -42,8 +45,7 @@ def moving_average(arr, window_size): """Calculates the moving average using NumPy's convolution function.""" # Pad with zeros to handle initial window positions arr_padded = np.pad(arr, (window_size - 1, 0), "constant") - smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / - window_size) + smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / window_size) return [smoothed_arr, arr - smoothed_arr] @@ -464,6 +466,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): From 373a5293df29c060040e71ecc39e6665af463df5 Mon Sep 17 00:00:00 2001 From: Killer3048 Date: Wed, 9 Apr 2025 13:21:04 +0700 Subject: [PATCH 2/5] change line for smoother_arr --- src/timesfm/timesfm_base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/timesfm/timesfm_base.py b/src/timesfm/timesfm_base.py index 4502d21..b0a622a 100644 --- a/src/timesfm/timesfm_base.py +++ b/src/timesfm/timesfm_base.py @@ -45,7 +45,8 @@ def moving_average(arr, window_size): """Calculates the moving average using NumPy's convolution function.""" # Pad with zeros to handle initial window positions arr_padded = np.pad(arr, (window_size - 1, 0), "constant") - smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / window_size) + smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / + window_size) return [smoothed_arr, arr - smoothed_arr] From 6678ba6cee45c487406e260a3b4167e44d3aa317 Mon Sep 17 00:00:00 2001 From: Killer3048 Date: Thu, 10 Apr 2025 09:47:07 +0700 Subject: [PATCH 3/5] update pyproject (exclude jax/pax) and update README --- README.md | 16 +++++++++++----- pyproject.toml | 2 +- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index c29d978..1dc169f 100644 --- a/README.md +++ b/README.md @@ -78,18 +78,24 @@ 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. -Please see the `experiments` section fro more instructions. +If you plan to use the **`forecast_with_covariates`** function (which requires external regressors), +you need to install **JAX** and **jaxlib**. Installing TimesFM with either the `[pax]` or the `[torch]` extras will include these packages by default. +However, if you installed the base version of TimesFM, you must manually install the dependencies: +``` +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) diff --git a/pyproject.toml b/pyproject.toml index d2084f0..09a0a18 100644 --- a/pyproject.toml +++ b/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" From b1c6f4d925ae2730b6296afa0d1c9f63985fbd53 Mon Sep 17 00:00:00 2001 From: Killer3048 Date: Thu, 10 Apr 2025 09:52:58 +0700 Subject: [PATCH 4/5] update readme.md --- README.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/README.md b/README.md index 1dc169f..31208e5 100644 --- a/README.md +++ b/README.md @@ -81,8 +81,7 @@ After than you can run the timesfm under `poetry shell` or do `poetry run python **Additional Note**: If you plan to use the **`forecast_with_covariates`** function (which requires external regressors), -you need to install **JAX** and **jaxlib**. Installing TimesFM with either the `[pax]` or the `[torch]` extras will include these packages by default. -However, if you installed the base version of TimesFM, you must manually install the dependencies: +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 ``` From 6da174d703aab0ec3c6ae31e821c3d92bea800d8 Mon Sep 17 00:00:00 2001 From: Killer3048 Date: Thu, 10 Apr 2025 10:47:44 +0700 Subject: [PATCH 5/5] add information about install pax/paxlib for covariate support in section `Covariates Support` --- README.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/README.md b/README.md index 31208e5..fbe0d9f 100644 --- a/README.md +++ b/README.md @@ -246,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).