Merge pull request #145 from google-research/rajat_dev

Full pytorch support

README.md

@@ -16,14 +16,14 @@ This is not an officially supported Google product.
 
 We recommend at least 16GB RAM to load TimesFM dependencies.
 
-## Update - Aug. 6, 2024
-
+## Update - Sep. 12, 2024
+- We have released full pytorch support (excluding PEFT parts).
 - Shoutout to @tanmayshishodia for checking in PEFT methods like LoRA and DoRA.
 - To install TimesFM, you can now simply do: `pip install timesfm`.
 - Launched [finetuning support](https://github.com/google-research/timesfm/blob/master/notebooks/finetuning.ipynb) that lets you finetune the weights of the pretrained TimesFM model on your own data.
 - Launched [~zero-shot covariate support](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb) with external regressors. More details [here](https://github.com/google-research/timesfm?tab=readme-ov-file#covariates-support).
 
-## Checkpoint timesfm-1.0-200m
+## Checkpoint timesfm-1.0-200m (-pytorch)
 
 timesfm-1.0-200m is the first open model checkpoint:
 
@@ -39,68 +39,55 @@ Please look into the README files in the respective benchmark directories within
 
 ## Installation
 
-### Installation as a package
+### Local installation using poetry
 
-To install the TimesFM as a package, you can run the following command without cloning this repo:
-
-`pip install timesfm`
-
-### Installation using conda
-
-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:
+We will be using `pyenv` and `poetry`. In order to set these things up please follow the instructions [here](https://substack.com/home/post/p-148747960?r=28a5lx&utm_campaign=post&utm_medium=web). Note that the PAX (or JAX) version needs to run on python 3.10.x and the PyTorch version can run on >=3.11.x. Therefore make sure you have two versions of python installed:
 
 ```
-conda env create --file=environment.yml
+pyenv install 3.10
+pyenv install 3.11
+pyenv versions # to list the versions available (lets assume the versions are 3.10.15 and 3.11.10)
 ```
 
-For a CPU setup please use,
+### For PAX version installation do the following.
 
 ```
-conda env create --file=environment_cpu.yml
+pyenv local 3.10.15
+poetry env use 3.10.15
+poetry lock
+poetry install --only  pax
 ```
-to create the environment instead.
 
-Follow by
+After than you can run the timesfm under `poetry shell` or do `poetry run python3 ...`.
+
+### For PyTorch version installation do the following.
 
 ```
-conda activate tfm_env
-pip install -e .
+pyenv local 3.11.10
+poetry env use 3.11.10
+poetry lock
+poetry install --only  torch
 ```
-to install the package.
+
+After than you can run the timesfm under `poetry shell` or do `poetry run python3 ...`.
 
 **Note**: 
 
 1. Running the provided benchmarks would require additional dependencies.
-Please use the environment files under `experiments` instead.
+Please see the `experiments` section fro more instructions.
 
 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.
 
-
-### Local installation using poetry
-
-To from the current repository/local version (like you would have previously done with `pip -e .`), you can run the command
-
-```
-pip install poetry # optional
-poetry install
-```
-
-This will install the environment in the local .venv folder (depends on the configuration) and matches the python command to the poetry environment. If this is not the case, you can use `poetry run python` to use the local environment.
-
 ### Notes
 
-1. Running the provided benchmarks would require additional dependencies.
-Please use the environment files under `experiments` instead.
+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 code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned.
+2. The dependency `lingvo` does not support ARM architectures, and the PAX version is not working for machines with Apple silicon.
 
-#### Building the package and publishing to PyPI
+### Install from PyPI (and publish)
 
-The package can be built using the command `poetry build`.
+Instructions coming soon.
 
-To build and publish it to PyPI, the command `poetry publish` can be used. This command will require the user to have the necessary permissions to publish to the PyPI repository.
 
 ## Usage 
 
@@ -110,32 +97,36 @@ Then the base class can be loaded as,
 ```python
 import timesfm
 
+# For PAX
 tfm = timesfm.TimesFm(
-    context_len=<context>,
-    horizon_len=<horizon>,
-    input_patch_len=32,
-    output_patch_len=128,
-    num_layers=20,
-    model_dims=1280,
-    backend=<backend>,
-)
-tfm.load_from_checkpoint(repo_id="google/timesfm-1.0-200m")
+      hparams=timesfm.TimesFmHparams(
+          backend="gpu",
+          per_core_batch_size=32,
+          horizon_len=128,
+      ),
+      checkpoint=timesfm.TimesFmCheckpoint(
+          huggingface_repo_id="google/timesfm-1.0-200m"),
+  )
+
+# For Torch
+tfm = timesfm.TimesFm(
+      hparams=timesfm.TimesFmHparams(
+          backend="gpu",
+          per_core_batch_size=32,
+          horizon_len=128,
+      ),
+      checkpoint=timesfm.TimesFmCheckpoint(
+          huggingface_repo_id="google/timesfm-1.0-200m-pytorch"),
+  )
 ```
 
-Note that the four parameters are fixed to load the 200m model
+Note some of the parameters are fixed to load the 200m model
 
-```python
-input_patch_len=32,
-output_patch_len=128,
-num_layers=20,
-model_dims=1280,
-```
-
-1. The `context_len` here can be set as the max context length **of the model**. **It needs to be a multiplier of `input_patch_len`, i.e. a multiplier of 32.** You can provide a shorter series to the `tfm.forecast()` function and the model will handle it. Currently, the model handles a max context length of 512, which can be increased in later releases. The input time series can have **any context length**. Padding / truncation will be handled by the inference code if needed.
+1. The `context_len` in `hparams` here can be set as the max context length **of the model**. **It needs to be a multiplier of `input_patch_len`, i.e. a multiplier of 32.** You can provide a shorter series to the `tfm.forecast()` function and the model will handle it. Currently, the model handles a max context length of 512, which can be increased in later releases. The input time series can have **any context length**. Padding / truncation will be handled by the inference code if needed.
 
 2. The horizon length can be set to anything. We recommend setting it to the largest horizon length you would need in the forecasting tasks for your application. We generally recommend horizon length <= context length but it is not a requirement in the function call.
 
-3. `backend` is one of "cpu", "gpu" or "tpu", case sensitive.
+3. `backend` is one of "cpu", "gpu", case sensitive.
 
 ### Perform inference
 

environment.yml

@@ -1,21 +0,0 @@
-name: tfm_env
-
-channels:
-  - conda-forge
-  - defaults
-  - anaconda
-dependencies:
-  - jupyterlab
-  - pip
-  - python=3.10
-  - pip:
-    - huggingface_hub[cli]
-    - utilsforecast
-    - praxis
-    - paxml
-    - jax[cuda12]==0.4.26
-    - einshape
-    - scikit-learn
-    - typer
-    - wandb
-    - pytest

environment_cpu.yml

@@ -1,21 +0,0 @@
-name: tfm_env
-
-channels:
-  - conda-forge
-  - defaults
-  - anaconda
-dependencies:
-  - jupyterlab
-  - pip
-  - python=3.10
-  - pip:
-    - huggingface_hub[cli]
-    - utilsforecast
-    - praxis
-    - paxml
-    - jax[cpu]==0.4.26
-    - einshape
-    - scikit-learn
-    - typer
-    - wandb
-    - pytest

experiments/environment.yml

@@ -1,28 +0,0 @@
-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[cuda12]==0.4.26
-    - einshape
-    - python-dotenv
-    - nixtla>=0.5.1
-    - rich
-    - scikit-learn

experiments/environment_cpu.yml

@@ -1,28 +0,0 @@
-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]==0.4.26
-    - einshape
-    - python-dotenv
-    - nixtla>=0.5.1
-    - rich
-    - scikit-learn

experiments/extended_benchmarks/README.md

@@ -5,14 +5,20 @@ The benchmark setting has been borrowed from Nixtla's original [benchmarking](ht
 
 ## Running TimesFM on the benchmark
 
-Install the environment and the package as detailed in the main README and then follow the steps from the base directory.
+We need to add the following packages for running these benchmarks. Follow the installation instructions till before `poetry lock`. Then,
 
 ```
-conda activate tfm_env
-TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false python3 -m experiments.extended_benchmarks.run_timesfm --model_path=<model_path> --backend="gpu"
+poetry add git+https://github.com/awslabs/gluon-ts.git
+poetry lock
+poetry install --only <pax or pytorch>
+```
+
+To run the timesfm on the benchmark do:
+
+```
+poetry run python3 -m experiments.extended_benchmarks.run_timesfm --model_path=google/timesfm-1.0-200m(-pytorch) --backend="gpu"
 ```
 
-In the above, `<model_path>` should point to the checkpoint directory that can be downloaded from HuggingFace. 
 
 Note: In the current version of TimesFM we focus on point forecasts and therefore the mase, smape have been calculated using the quantile head corresponding to the median i.e 0.5 quantile. We do offer 10 quantile heads but they have not been calibrated after pretraining. We recommend using them with caution or calibrate/conformalize them on a hold out for your applications. More to follow on later versions.
 
@@ -22,7 +28,7 @@ Note: In the current version of TimesFM we focus on point forecasts and therefor
 
 __Update:__ We have added TimeGPT-1 to the benchmark results. We had to remove the Dominick dataset as we were not able to run TimeGPT-1 on this benchmark. Note that the previous results including Dominick remain available at `./tfm_results.png`. In order to reproduce the results for TimeGPT-1, please run `run_timegpt.py`.
 
-_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT.
+_Remark:_ All baselines except the ones involving TimeGPT were run performed on a [g2-standard-32](https://cloud.google.com/compute/docs/gpus). Since TimeGPT-1 can only be accessed by an API, the time column might not reflect the true speed of the model as it also includes the communication cost. Moreover, we are not sure about the exact backend hardware for TimeGPT. The TimesFM latency numbers are from the PAX version.
 
 We can see that TimesFM performs the best in terms of both mase and smape. More importantly it is much faster than the other methods, in particular it is more than 600x faster than StatisticalEnsemble and 80x faster than Chronos (Large).
 

experiments/extended_benchmarks/run_timesfm.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Evaluation script for timesfm."""
 
 import os
@@ -21,12 +20,10 @@ import time
 from absl import flags
 import numpy as np
 import pandas as pd
-from paxml import checkpoints
 import timesfm
 
 from .utils import ExperimentHandler
 
-
 dataset_names = [
     "m1_monthly",
     "m1_quarterly",
@@ -74,35 +71,27 @@ context_dict = {
     "m4_yearly": 64,
 }
 
-_MODEL_PATH = flags.DEFINE_string(
-    "model_path", "/home/timesfm_q10_20240501", "Path to model"
-)
+_MODEL_PATH = flags.DEFINE_string("model_path", "google/timesfm-1.0-200m",
+                                  "Path to model")
 _BATCH_SIZE = flags.DEFINE_integer("batch_size", 64, "Batch size")
 _HORIZON = flags.DEFINE_integer("horizon", 128, "Horizon")
 _BACKEND = flags.DEFINE_string("backend", "gpu", "Backend")
 _NUM_JOBS = flags.DEFINE_integer("num_jobs", 1, "Number of jobs")
 _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
 
-
 QUANTILES = list(np.arange(1, 10) / 10.0)
 
 
 def main():
   results_list = []
   tfm = timesfm.TimesFm(
-      context_len=512,
-      horizon_len=_HORIZON.value,
-      input_patch_len=32,
-      output_patch_len=128,
-      num_layers=20,
-      model_dims=1280,
-      backend=_BACKEND.value,
-      per_core_batch_size=_BATCH_SIZE.value,
-      quantiles=QUANTILES,
-  )
-  tfm.load_from_checkpoint(
-      _MODEL_PATH.value,
-      checkpoint_type=checkpoints.CheckpointType.FLAX,
+      hparams=timesfm.TimesFmHparams(
+          backend=_BACKEND.value,
+          per_core_batch_size=_BATCH_SIZE.value,
+          horizon_len=_HORIZON.value,
+      ),
+      checkpoint=timesfm.TimesFmCheckpoint(
+          huggingface_repo_id=_MODEL_PATH.value),
   )
   run_id = np.random.randint(100000)
   model_name = "timesfm"
@@ -127,9 +116,9 @@ def main():
     )
     total_time = time.time() - init_time
     time_df = pd.DataFrame({"time": [total_time], "model": model_name})
-    results = exp.evaluate_from_predictions(
-        models=[model_name], fcsts_df=fcsts_df, times_df=time_df
-    )
+    results = exp.evaluate_from_predictions(models=[model_name],
+                                            fcsts_df=fcsts_df,
+                                            times_df=time_df)
     print(results, flush=True)
     results_list.append(results)
     results_full = pd.concat(results_list)

experiments/long_horizon_benchmarks/README.md

@@ -6,12 +6,21 @@ All experiments were performed on a [g2-standard-32](https://cloud.google.com/co
 
 ## Running TimesFM on the benchmark
 
-Install the environment and the package as detailed in the main README and then follow the steps from the base directory.
+We need to add the following packages for running these benchmarks. Follow the installation instructions till before `poetry lock`. Then,
 
 ```
-conda activate tfm_env
-TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false python3 -m experiments.long_horizon_benchmarks.run_eval \
---model_path=<model_path> --backend="gpu" \
+poetry add git+https://github.com/awslabs/gluon-ts.git
+poetry add git+https://github.com/amazon-science/chronos-forecasting.git
+poetry lock
+poetry install --only pax
+```
+Note that for now only the pax version runs on this benchmark, because we had to remove the old tf dependency from the pytorch version. We will fix this issue soon.
+
+To run the timesfm on the benchmark do:
+
+```
+poetry run python3 -m experiments.long_horizon_benchmarks.run_eval \
+--model_path=google/timesfm-1.0-200m --backend="gpu" \
 --pred_len=96 --context_len=512 --dataset=etth1
 ```
 
@@ -20,7 +29,7 @@ In the above, `<model_path>` should point to the checkpoint directory that can b
 For running chronos on the same benchmark you can run the command,
 
 ```
-TF_CPP_MIN_LOG_LEVEL=2 XLA_PYTHON_CLIENT_PREALLOCATE=false python3 -m experiments.long_horizon_benchmarks.run_eval \
+poetry run python3 -m experiments.long_horizon_benchmarks.run_eval \
 --model_path=amazon/chronos-t5-mini --backend="gpu" \
 --pred_len=96 --context_len=512 --dataset=etth1
 ```

experiments/long_horizon_benchmarks/run_eval.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Eval pipeline."""
 
 import json
@@ -22,44 +21,34 @@ from absl import flags
 import chronos
 import numpy as np
 import pandas as pd
-from paxml import checkpoints
 import timesfm
 from timesfm import data_loader
 import torch
 import tqdm
 
-
 FLAGS = flags.FLAGS
 
-_BATCH_SIZE = flags.DEFINE_integer(
-    "batch_size", 64, "Batch size for the randomly sampled batch"
-)
+_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64,
+                                   "Batch size for the randomly sampled batch")
 _DATASET = flags.DEFINE_string("dataset", "etth1", "The name of the dataset.")
-_MODEL_PATH = flags.DEFINE_string(
-    "model_path", "./timesfm_q10_20240501", "Path to model."
-)
-_DATETIME_COL = flags.DEFINE_string(
-    "datetime_col", "date", "Column having datetime."
-)
-_NUM_COV_COLS = flags.DEFINE_list(
-    "num_cov_cols", None, "Column having numerical features."
-)
-_CAT_COV_COLS = flags.DEFINE_list(
-    "cat_cov_cols", None, "Column having categorical features."
-)
+
+_MODEL_PATH = flags.DEFINE_string("model_path", "./timesfm_q10_20240501",
+                                  "The name of the dataset.")
+_DATETIME_COL = flags.DEFINE_string("datetime_col", "date",
+                                    "Column having datetime.")
+_NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None,
+                                  "Column having numerical features.")
+_CAT_COV_COLS = flags.DEFINE_list("cat_cov_cols", None,
+                                  "Column having categorical features.")
 _TS_COLS = flags.DEFINE_list("ts_cols", None, "Columns of time-series features")
-_NORMALIZE = flags.DEFINE_bool(
-    "normalize", True, "normalize data for eval or not"
-)
-_CONTEXT_LEN = flags.DEFINE_integer(
-    "context_len", 512, "Length of the context window"
-)
+_NORMALIZE = flags.DEFINE_bool("normalize", True,
+                               "normalize data for eval or not")
+_CONTEXT_LEN = flags.DEFINE_integer("context_len", 512,
+                                    "Length of the context window")
 _PRED_LEN = flags.DEFINE_integer("pred_len", 96, "prediction length.")
 _BACKEND = flags.DEFINE_string("backend", "gpu", "backend to use")
-_RESULTS_DIR = flags.DEFINE_string(
-    "results_dir", "./results/long_horizon", "results directory"
-)
-
+_RESULTS_DIR = flags.DEFINE_string("results_dir", "./results/long_horizon",
+                                   "results directory")
 
 DATA_DICT = {
     "ettm2": {
@@ -176,9 +165,8 @@ def eval():
       holiday=False,
       permute=False,
   )
-  eval_itr = dtl.tf_dataset(
-      mode="test", shift=_PRED_LEN.value
-  ).as_numpy_iterator()
+  eval_itr = dtl.tf_dataset(mode="test",
+                            shift=_PRED_LEN.value).as_numpy_iterator()
   model_path = _MODEL_PATH.value
   if model_path.startswith("amazon"):
     model = chronos.ChronosPipeline.from_pretrained(
@@ -188,19 +176,12 @@ def eval():
     )
   else:
     model = timesfm.TimesFm(
-        context_len=_CONTEXT_LEN.value,
-        horizon_len=_PRED_LEN.value,
-        input_patch_len=32,
-        output_patch_len=128,
-        num_layers=20,
-        model_dims=1280,
-        backend=_BACKEND.value,
-        per_core_batch_size=batch_size,
-        quantiles=QUANTILES,
-    )
-    model.load_from_checkpoint(
-        model_path,
-        checkpoint_type=checkpoints.CheckpointType.FLAX,
+        hparams=timesfm.TimesFmHparams(
+            backend=_BACKEND.value,
+            per_core_batch_size=_BATCH_SIZE.value,
+            horizon_len=_PRED_LEN.value,
+        ),
+        checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=model_path),
     )
   smape_run_losses = []
   mse_run_losses = []
@@ -213,10 +194,9 @@ def eval():
   for batch in tqdm.tqdm(eval_itr):
     past = batch[0]
     actuals = batch[3]
-    forecasts = get_forecasts(
-        model_path, model, past, int_freq, _PRED_LEN.value
-    )
-    forecasts = forecasts[:, 0 : actuals.shape[1]]
+    forecasts = get_forecasts(model_path, model, past, int_freq,
+                              _PRED_LEN.value)
+    forecasts = forecasts[:, 0:actuals.shape[1]]
     mae_run_losses.append(_mae(forecasts, actuals).sum())
     mse_run_losses.append(_mse(forecasts, actuals).sum())
     smape_run_losses.append(_smape(forecasts, actuals).sum())

pyproject.toml

@@ -1,8 +1,6 @@
 [tool.poetry]
 name = "timesfm"
-packages = [
-    { include = "timesfm", from = "src" },
-]
+packages = [{ include = "timesfm", from = "src" }]
 description = "Open weights time-series foundation model from Google Research."
 version = "1.0.1"
 authors = [
@@ -24,27 +22,41 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
     "Topic :: Software Development :: Quality Assurance",
 ]
-include = [
-    "LICENSE",
-]
+include = ["LICENSE"]
 
 [tool.poetry.dependencies]
-python = ">=3.10,<3.11"
+python = ">=3.10,<3.12"
 einshape = ">=1.0.0"
 numpy = ">=1.26.4"
 pandas = ">=2.1.4"
-paxml = ">=1.4.0"
 utilsforecast = ">=0.1.10"
-jax = {version = ">=0.4.26", extras = ["cuda12"]}
-jaxlib = ">=0.4.26"
-huggingface_hub = {version = ">=0.23.0", extras = ["cli"]}
+huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
 scikit-learn = ">=1.2.2"
 typer = ">=0.12.3"
 wandb = ">=0.17.5"
+ipython = "^8.27.0"
+absl-py = ">=1.4.0"
+
+[tool.poetry.group.pax]
+optional = true
+
+[tool.poetry.group.pax.dependencies]
+paxml = { version = ">=1.4.0", python = ">=3.10,<3.11" }
+lingvo = { version = ">=0.12.7", python = ">=3.10,<3.11" }
+jax = { version = ">=0.4.26", extras = ["cuda12"], python = ">=3.10,<3.11" }
+jaxlib = { version = ">=0.4.26", python = ">=3.10,<3.11" }
+
+[tool.poetry.group.torch]
+optional = true
+
+[tool.poetry.group.torch.dependencies]
+torch = { version = ">=2.0.0", extras = ["cuda"], python = ">=3.11,<3.12" }
+jax = { version = ">=0.4.26", extras = ["cuda12"], python = ">=3.11,<3.12" }
+jaxlib = { version = ">=0.4.26", python = ">=3.11,<3.12" }
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=8.3.2"
 
 [build-system]
 requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+build-backend = "poetry.core.masonry.api"

src/timesfm/__init__.py

@@ -11,7 +11,15 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """TimesFM init file."""
-
-from .timesfm import TimesFm, freq_map
+print(
+    "TimesFM v1.2.0. See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs."
+)
+from timesfm.timesfm_base import freq_map, TimesFmCheckpoint, TimesFmHparams, TimesFmBase
+try:
+  print("Loaded Jax TimesFM.")
+  from timesfm.timesfm_jax import TimesFmJax as TimesFm
+  from timesfm import data_loader
+except Exception as _:
+  print("Loaded PyTorch TimesFM.")
+  from timesfm.timesfm_torch import TimesFmTorch as TimesFm

src/timesfm/patched_decoder.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Pax ML model for patched time-series decoder.
 
 The file implements Residual MLPs, Patched Decoder layers and PAX ML models.
@@ -36,7 +35,6 @@ from praxis.layers import normalizations
 from praxis.layers import stochastics
 from praxis.layers import transformers
 
-
 # PAX shortcuts
 NestedMap = py_utils.NestedMap
 JTensor = pytypes.JTensor
@@ -44,7 +42,6 @@ JTensor = pytypes.JTensor
 LayerTpl = pax_fiddle.Config[base_layer.BaseLayer]
 template_field = base_layer.template_field
 
-
 PAD_VAL = 1123581321.0
 DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
 
@@ -57,7 +54,6 @@ _FREQ = "freq"
 _OUTPUT_TOKENS = "output_tokens"
 _STATS = "stats"
 
-
 # Small numerical value.
 _TOLERANCE = 1e-7
 
@@ -158,9 +154,8 @@ class ResidualBlock(base_layer.BaseLayer):
       return output + residual
 
 
-def _masked_mean_std(
-    inputs: JTensor, padding: JTensor
-) -> Tuple[JTensor, JTensor]:
+def _masked_mean_std(inputs: JTensor,
+                     padding: JTensor) -> Tuple[JTensor, JTensor]:
   """Calculates mean and standard deviation of arr across axis 1.
 
   It should exclude values where pad is 1.
@@ -197,7 +192,7 @@ def _masked_mean_std(
 
   # Calculate the masked sum and squared sum of M
   masked_sum = jnp.sum(arr * mask, axis=1)
-  masked_squared_sum = jnp.sum((arr * mask) ** 2, axis=1)
+  masked_squared_sum = jnp.sum((arr * mask)**2, axis=1)
 
   # Calculate the masked mean and standard deviation
   masked_mean = masked_sum / num_valid_elements
@@ -240,8 +235,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
   quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
   residual_block_tpl: LayerTpl = template_field(ResidualBlock)
   stacked_transformer_params_tpl: LayerTpl = template_field(
-      transformers.StackedTransformer
-  )
+      transformers.StackedTransformer)
   use_freq: bool = True
 
   def setup(self) -> None:
@@ -276,9 +270,8 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
 
     self.create_child(
         "position_emb",
-        pax_fiddle.Config(
-            layers.PositionalEmbedding, embedding_dims=self.model_dims
-        ),
+        pax_fiddle.Config(layers.PositionalEmbedding,
+                          embedding_dims=self.model_dims),
     )
 
     if self.use_freq:
@@ -292,27 +285,24 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
       )
 
   def transform_decode_state(
-      self, transform_fn: base_layer.DecodeStateTransformFn
-  ) -> None:
+      self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
     """Transforms all decode state variables based on transform_fn."""
     self.stacked_transformer_layer.transform_decode_state(transform_fn)
 
   def _forward_transform(
-      self, inputs: JTensor, patched_pads: JTensor
-  ) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
+      self, inputs: JTensor,
+      patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
     """Input is of shape [B, N, P]."""
     mu, sigma = _masked_mean_std(inputs, patched_pads)
     sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
     # Normalize each patch.
     outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
     outputs = jnp.where(
-        jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs
-    )
+        jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs)
     return outputs, (mu, sigma)
 
-  def _reverse_transform(
-      self, outputs: JTensor, stats: Tuple[JTensor, JTensor]
-  ) -> JTensor:
+  def _reverse_transform(self, outputs: JTensor,
+                         stats: Tuple[JTensor, JTensor]) -> JTensor:
     """Output is of shape [B, N, P, Q]."""
     mu, sigma = stats
     return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
@@ -326,18 +316,15 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     """Preprocess input for stacked transformer."""
     # Reshape into patches.
     patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len)
-    patched_pads = es.jax_einshape(
-        "b(np)->bnp", input_padding, p=self.patch_len
-    )
+    patched_pads = es.jax_einshape("b(np)->bnp",
+                                   input_padding,
+                                   p=self.patch_len)
     patched_inputs = jnp.where(
-        jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs
-    )
+        jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs)
     patched_pads = jnp.where(
-        jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads
-    )
-    patched_inputs, stats = self._forward_transform(
-        patched_inputs, patched_pads
-    )
+        jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads)
+    patched_inputs, stats = self._forward_transform(patched_inputs,
+                                                    patched_pads)
 
     # B x N x D
     patched_inputs = patched_inputs * (1.0 - patched_pads)
@@ -367,9 +354,10 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     """Postprocess output of stacked transformer."""
     # B x N x (H.Q)
     output_ts = self.horizon_ff_layer(model_output)
-    output_ts = es.jax_einshape(
-        "bn(hq)->bnhq", output_ts, q=num_outputs, h=self.horizon_len
-    )
+    output_ts = es.jax_einshape("bn(hq)->bnhq",
+                                output_ts,
+                                q=num_outputs,
+                                h=self.horizon_len)
     return self._reverse_transform(output_ts, stats)
 
   def __call__(self, inputs: NestedMap) -> NestedMap:
@@ -400,9 +388,11 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     model_output = self.stacked_transformer_layer(model_input, patched_padding)
 
     output_ts = self._postprocess_output(model_output, num_outputs, stats)
-    return NestedMap(
-        {_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats}
-    )
+    return NestedMap({
+        _OUTPUT_TOKENS: model_output,
+        _OUTPUT_TS: output_ts,
+        _STATS: stats
+    })
 
   def decode(
       self,
@@ -443,15 +433,13 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     if paddings.shape[1] != final_out.shape[1] + horizon_len:
       raise ValueError(
           "Length of paddings must match length of input + horizon_len:"
-          f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}"
-      )
+          f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}")
     if output_patch_len is None:
       output_patch_len = self.horizon_len
-    num_decode_patches = (
-        horizon_len + output_patch_len - 1
-    ) // output_patch_len
+    num_decode_patches = (horizon_len + output_patch_len -
+                          1) // output_patch_len
     for step_index in range(num_decode_patches):
-      current_padding = paddings[:, 0 : final_out.shape[1]]
+      current_padding = paddings[:, 0:final_out.shape[1]]
       input_ts = final_out[:, -max_len:]
       input_padding = current_padding[:, -max_len:]
       model_input = NestedMap(
@@ -463,7 +451,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
       if return_forecast_on_context and step_index == 0:
         # For the first decodings step, collect the model forecast on the
         # context except the unavailable first input batch forecast.
-        new_full_ts = fprop_outputs[:, :-1, : self.patch_len, :]
+        new_full_ts = fprop_outputs[:, :-1, :self.patch_len, :]
         new_full_ts = es.jax_einshape("bnph->b(np)h", new_full_ts)
 
         full_outputs.append(new_full_ts)
@@ -477,9 +465,9 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
 
     if return_forecast_on_context:
       # `full_outputs` indexing starts at after the first input patch.
-      full_outputs = jnp.concatenate(full_outputs, axis=1)[
-          :, : (context_len - self.patch_len + horizon_len), :
-      ]
+      full_outputs = jnp.concatenate(full_outputs,
+                                     axis=1)[:, :(context_len - self.patch_len +
+                                                  horizon_len), :]
     else:
       # `full_outputs` indexing starts at the forecast horizon.
       full_outputs = jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :]
@@ -506,14 +494,12 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
     input_padding = jnp.zeros_like(input_ts)
     context_len = input_ts.shape[1]
     input_patch_len = self.core_layer_tpl.patch_len
-    context_pad = (
-        (context_len + input_patch_len - 1) // input_patch_len
-    ) * input_patch_len - context_len
+    context_pad = ((context_len + input_patch_len - 1) //
+                   input_patch_len) * input_patch_len - context_len
 
     input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
-    input_padding = jnp.pad(
-        input_padding, [(0, 0), (context_pad, 0)], constant_values=1
-    )
+    input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)],
+                            constant_values=1)
     freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
     new_input_batch = NestedMap(
         input_ts=input_ts,
@@ -522,9 +508,8 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
     )
     return self.core_layer(new_input_batch)
 
-  def _quantile_loss(
-      self, pred: JTensor, actual: JTensor, quantile: float
-  ) -> JTensor:
+  def _quantile_loss(self, pred: JTensor, actual: JTensor,
+                     quantile: float) -> JTensor:
     """Calculates quantile loss.
 
     Args:
@@ -540,12 +525,11 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
     loss_second = -dev * (1.0 - quantile)
     return 2 * jnp.where(loss_first >= 0, loss_first, loss_second)
 
-  def compute_loss(
-      self, prediction_output: NestedMap, input_batch: NestedMap
-  ) -> Tuple[NestedMap, NestedMap]:
+  def compute_loss(self, prediction_output: NestedMap,
+                   input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]:
     output_ts = prediction_output[_OUTPUT_TS]
     actual_ts = input_batch[_TARGET_FUTURE]
-    pred_ts = output_ts[:, -1, 0 : actual_ts.shape[1], :]
+    pred_ts = output_ts[:, -1, 0:actual_ts.shape[1], :]
     loss = jnp.square(pred_ts[:, :, 0] - actual_ts)
     for i, quantile in enumerate(self.core_layer.quantiles):
       loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile)

src/timesfm/pytorch_patched_decoder.py

@@ -22,7 +22,7 @@ import torch.nn.functional as F
 
 
 def _create_quantiles() -> list[float]:
-  return [0.1, 0.25, 0.5, 0.75, 0.9]
+  return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
 
 
 @dataclasses.dataclass
@@ -192,7 +192,8 @@ def convert_paddings_to_mask(
   Returns:
       A torch.Tensor of shape [B, 1, 1, T] ready to add to attention logits.
   """
-  attention_mask = paddings[:, None, None, :]  # Equivalent to jnp.newaxis
+  attention_mask = paddings.detach().clone()
+  attention_mask = attention_mask[:, None, None, :]  # Equivalent to jnp.newaxis
   attention_mask *= get_large_negative_number(dtype)
   return attention_mask
 

src/timesfm/timesfm_base.py

@@ -11,43 +11,26 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""TimesFM forecast API for inference."""
+"""Base class for TimesFM inference. This will be common to PAX and Pytorch."""
 
 import collections
+import dataclasses
 import logging
 import multiprocessing
-from os import path
-import time
-from typing import Any, Literal, Optional, Sequence
+from typing import Any, Literal, Sequence
 
-import einshape as es
-from huggingface_hub import snapshot_download
-import jax
-import jax.numpy as jnp
 import numpy as np
 import pandas as pd
-from paxml import checkpoints
-from paxml import tasks_lib
-from praxis import base_hyperparams
-from praxis import base_layer
-from praxis import pax_fiddle
-from praxis import py_utils
-from praxis import pytypes
-from praxis.layers import normalizations
-from praxis.layers import transformers
 
 from utilsforecast.processing import make_future_dataframe
 
-from . import patched_decoder
 from . import xreg_lib
 
-instantiate = base_hyperparams.instantiate
-NestedMap = py_utils.NestedMap
-JTensor = pytypes.JTensor
 Category = xreg_lib.Category
 XRegMode = xreg_lib.XRegMode
 
 _TOL = 1e-6
+DEFAULT_QUANTILES = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
 
 
 def process_group(key, group, value_name, forecast_context_len):
@@ -79,7 +62,7 @@ def freq_map(freq: str):
 
 
 # Per time series normalization: forward.
-def _normalize(batch):
+def normalize(batch):
   stats = [
       (np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch
   ]
@@ -88,32 +71,19 @@ def _normalize(batch):
 
 
 # Per time series normalization: inverse.
-def _renormalize(batch, stats):
+def renormalize(batch, stats):
   return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)]
 
 
-class TimesFm:
-  """TimesFM forecast API for inference.
+@dataclasses.dataclass(kw_only=True)
+class TimesFmHparams:
+  """Hparams used to initialize a TimesFM model for inference.
 
-  This class is the scaffolding for calling TimesFM forecast. To properly use:
-    1. Create an instance with the correct hyperparameters of a TimesFM model.
-    2. Call `load_from_checkpoint` to load a compatible checkpoint.
-    3. Call `forecast` for inference.
-
-  Given the model size, this API does not shard the model weights for SPMD. All
-  parallelism happens on the data dimension.
-
-  Compilation happens during the first time `forecast` is called and uses the
-  `per_core_batch_size` to set and freeze the input signature. Subsequent calls
-  to `forecast` reflect the actual inference latency.
+  These are the sufficient subset of hparams to configure TimesFM inference
+  agnostic to the checkpoint version, and are not necessarily the same as the
+  hparams used to train the checkpoint.
 
   Attributes:
-    per_core_batch_size: Batch size on each core for data parallelism.
-    backend: One of "cpu", "gpu" or "tpu".
-    num_devices: Number of cores provided the backend.
-    global_batch_size: per_core_batch_size * num_devices. Each batch of
-      inference task will be padded with respect to global_batch_size to
-      minimize latency.
     context_len: Largest context length the model allows for each decode call.
       This technically can be any large, but practically should set to the
       context length the checkpoint was trained with.
@@ -122,237 +92,96 @@ class TimesFm:
     output_patch_len: Output patch len. How many timepoints is taken from a
       single step of autoregressive decoding. Can be set as the training horizon
       of the checkpoint.
-    mesh_shape: Shape of the data parallelism mesh.
-    mesh_name: Names of the data parallelism mesh.
-    model_p: Configuration of the TimesFM model deduced from the hparams.
+    num_layers: Number of transformer layers in the model.
+    model_dims: Model dimension.
+    per_core_batch_size: Batch size on each core for data parallelism.
+    backend: One of "cpu", "gpu" or "tpu".
+    quantiles: Which quantiles are output by the model.
+  """
+
+  context_len: int = 512
+  horizon_len: int = 128
+  input_patch_len: int = 32
+  output_patch_len: int = 128
+  num_layers: int = 20
+  num_heads: int = 16
+  model_dims: int = 1280
+  per_core_batch_size: int = 32
+  backend: Literal["cpu", "gpu", "tpu"] = "cpu"
+  quantiles: Sequence[float] | None = DEFAULT_QUANTILES
+
+
+@dataclasses.dataclass(kw_only=True)
+class TimesFmCheckpoint:
+  """Checkpoint used to initialize a TimesFM model for inference.
+
+  Attributes:
+    version: Version of the checkpoint, e.g. "jax", "torch", "tensorflow", etc.
+      The factory will create the corresponding TimesFm inference class based on
+      this version.
+    path: Path to the checkpoint.
+    type: If provided, type of the checkpoint used by the specific checkpoint
+      loader per version.
+    step: If provided, step of the checkpoint.
+  """
+
+  version: str = "jax"
+  path: str | None = None
+  huggingface_repo_id: str | None = None
+  type: Any = None
+  step: int | None = None
+
+
+class TimesFmBase:
+  """Base TimesFM forecast API for inference.
+
+  This class is the scaffolding for calling TimesFM forecast. To properly use:
+    1. Create an instance with the correct hyperparameters of a TimesFM model.
+    2. Call `load_from_checkpoint` to load a compatible checkpoint.
+    3. Call `forecast` for inference.
   """
 
   def _logging(self, s):
-    if self._verbose:
-      print(s)
+    print(s)
 
-  def __init__(
-      self,
-      context_len: int,
-      horizon_len: int,
-      input_patch_len: int,
-      output_patch_len: int,
-      num_layers: int,
-      model_dims: int,
-      per_core_batch_size: int = 32,
-      backend: Literal["cpu", "gpu", "tpu"] = "cpu",
-      quantiles: Sequence[float] | None = None,
-      verbose: bool = True,
-  ) -> None:
+  def __post_init__(self) -> None:
+    """Additional initialization for subclasses before checkpoint loading."""
+    pass
+
+  def __init__(self, hparams: TimesFmHparams,
+               checkpoint: TimesFmCheckpoint) -> None:
     """Initializes the TimesFM forecast API.
 
     Args:
-      context_len: Largest context length the model allows for each decode call.
-        This technically can be any large, but practically should set to the
-        context length the checkpoint was trained with.
-      horizon_len: Forecast horizon.
-      input_patch_len: Input patch len.
-      output_patch_len: Output patch len. How many timepoints is taken from a
-        single step of autoregressive decoding. Can be set as the training
-        horizon of the checkpoint.
-      num_layers: Number of transformer layers.
-      model_dims: Model dimension.
-      per_core_batch_size: Batch size on each core for data parallelism.
-      backend: One of "cpu", "gpu" or "tpu".
-      quantiles: list of output quantiles supported by the model.
-      verbose: Whether to print logging messages.
+      hparams: Hyperparameters of the model.
+      checkpoint: Checkpoint to load. Notice `checkpoint.version` will decide
+        which TimesFM version to use.
     """
-    self.per_core_batch_size = per_core_batch_size
-    self.backend = backend
-    self.num_devices = jax.local_device_count(self.backend)
-    self.global_batch_size = self.per_core_batch_size * self.num_devices
+    self.hparams = hparams
+
+    # Expand hparams for conciseness within the model code.
+    self.context_len = hparams.context_len
+    self.horizon_len = hparams.horizon_len
+    self.input_patch_len = hparams.input_patch_len
+    self.output_patch_len = hparams.output_patch_len
+    self.num_layers = hparams.num_layers
+    self.model_dims = hparams.model_dims
+    self.backend = hparams.backend
+    self.quantiles = hparams.quantiles
+    self.num_heads = hparams.num_heads
+
+    # Rewrite these values in __post_init__ for SPMD.
+    self.num_cores = 1
+    self.per_core_batch_size = hparams.per_core_batch_size
+    self.global_batch_size = hparams.per_core_batch_size
 
-    self.context_len = context_len
-    self.horizon_len = horizon_len
-    self.input_patch_len = input_patch_len
-    self.output_patch_len = output_patch_len
     self._horizon_start = self.context_len - self.input_patch_len
+    self.__post_init__()
+    self.load_from_checkpoint(checkpoint)
 
-    self.mesh_shape = [1, self.num_devices, 1]
-    self.mesh_name = ["replica", "data", "mdl"]
-    if quantiles is None:
-      quantiles = patched_decoder.DEFAULT_QUANTILES
-
-    self.model_p = pax_fiddle.Config(
-        patched_decoder.PatchedTimeSeriesDecoder,
-        name="patched_decoder",
-        horizon_len=self.output_patch_len,
-        patch_len=input_patch_len,
-        model_dims=model_dims,
-        hidden_dims=model_dims,
-        residual_block_tpl=pax_fiddle.Config(patched_decoder.ResidualBlock),
-        quantiles=quantiles,
-        use_freq=True,
-        stacked_transformer_params_tpl=pax_fiddle.Config(
-            transformers.StackedTransformer,
-            num_heads=16,
-            num_layers=num_layers,
-            transformer_layer_params_tpl=pax_fiddle.Config(
-                transformers.Transformer,
-                ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
-            ),
-        ),
-    )
-
-    self._key1, self._key2 = jax.random.split(jax.random.PRNGKey(42))
-    self._model = None
-    self._train_state = None
-    self._pmapped_decode = None
-    self._verbose = verbose
-    self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
-    try:
-      multiprocessing.set_start_method("spawn")
-    except RuntimeError:
-      print("Multiprocessing context has already been set.")
-
-  def _get_sample_inputs(self):
-    return {
-        "input_ts":
-            jnp.zeros(
-                (
-                    self.per_core_batch_size,
-                    self.context_len + self.output_patch_len,
-                ),
-                dtype=jnp.float32,
-            ),
-        "input_padding":
-            jnp.zeros(
-                (
-                    self.per_core_batch_size,
-                    self.context_len + self.output_patch_len,
-                ),
-                dtype=jnp.float32,
-            ),
-        "freq":
-            jnp.zeros(
-                (
-                    self.per_core_batch_size,
-                    1,
-                ),
-                dtype=jnp.int32,
-            ),
-    }
-
-  def load_from_checkpoint(
-      self,
-      checkpoint_path: Optional[str] = None,
-      repo_id: str = "google/timesfm-1.0-200m",
-      checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.
-      FLAX,
-      step: int | None = None,
-  ) -> None:
-    """Loads a checkpoint and compiles the decoder.
-
-    Args:
-      checkpoint_path: Optional path to the checkpoint directory.
-      repo_id: Hugging Face Hub repo id.
-      checkpoint_type: type of PAX checkpoint
-      step: step of the checkpoint to load. If `None`, load latest checkpoint.
-    """
-    # Download the checkpoint from Hugging Face Hub if not given
-    if checkpoint_path is None:
-      checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints")
-
-    #  Initialize the model weights.
-    self._logging("Constructing model weights.")
-    start_time = time.time()
-    self._model = instantiate(self.model_p)
-    var_weight_hparams = self._model.abstract_init_with_metadata(
-        self._get_sample_inputs(), do_eval=True)
-    train_state_partition_specs = tasks_lib.create_state_partition_specs(
-        var_weight_hparams,
-        mesh_shape=self.mesh_shape,
-        mesh_axis_names=self.mesh_name,
-        discard_opt_states=True,
-        learners=None,
-    )
-    train_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
-        var_weight_hparams,
-        discard_opt_states=True,
-        learners=None,
-    )
-    self._logging(
-        f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
-
-    # Load the model weights.
-    self._logging(f"Restoring checkpoint from {checkpoint_path}.")
-    start_time = time.time()
-    self._train_state = checkpoints.restore_checkpoint(
-        train_state_local_shapes,
-        checkpoint_dir=checkpoint_path,
-        checkpoint_type=checkpoint_type,
-        state_specs=train_state_partition_specs,
-        step=step,
-    )
-    self._logging(
-        f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
-    self.jit_decode()
-
-  def jit_decode(self):
-    """Jitting decoding function."""
-
-    # Initialize and jit the decode fn.
-    def _decode(inputs):
-      assert self._model is not None
-      assert self._train_state is not None
-      return self._model.apply(
-          self._train_state.mdl_vars,
-          inputs,
-          horizon_len=self.horizon_len,
-          output_patch_len=self.output_patch_len,
-          max_len=self.context_len,
-          return_forecast_on_context=True,
-          rngs={
-              base_layer.PARAMS: self._key1,
-              base_layer.RANDOM: self._key2,
-          },
-          method=self._model.decode,
-      )
-
-    self._logging("Jitting decoding.")
-    start_time = time.time()
-    self._pmapped_decode = jax.pmap(
-        _decode,
-        axis_name="batch",
-        devices=jax.devices(self.backend),
-        backend=self.backend,
-        axis_size=self.num_devices,
-    )
-    with base_layer.JaxContext.new_context(hparams=self._eval_context):
-      _ = self._pmapped_decode(
-          NestedMap({
-              "input_ts":
-                  jnp.zeros(
-                      (
-                          self.num_devices,
-                          self.per_core_batch_size,
-                          self.context_len,
-                      ),
-                      dtype=jnp.float32,
-                  ),
-              "input_padding":
-                  jnp.zeros(
-                      (
-                          self.num_devices,
-                          self.per_core_batch_size,
-                          self.context_len + self.horizon_len,
-                      ),
-                      dtype=jnp.float32,
-                  ),
-              "date_features":
-                  None,
-              "freq":
-                  jnp.zeros(
-                      (self.num_devices, self.per_core_batch_size, 1),
-                      dtype=jnp.int32,
-                  ),
-          }))
-    self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
+  def load_from_checkpoint(self, checkpoint: TimesFmCheckpoint) -> None:
+    """Loads a checkpoint and compiles the decoder."""
+    raise NotImplementedError("`load_from_checkpoint` is not implemented.")
 
   def _preprocess(self, inputs: Sequence[np.array],
                   freq: Sequence[int]) -> tuple[np.array, np.array, int]:
@@ -417,7 +246,7 @@ class TimesFm:
       forecast_context_len: int | None = None,
       return_forecast_on_context: bool = False,
       truncate_negative: bool = False,
-  ) -> tuple[JTensor, JTensor]:
+  ) -> tuple[np.array, np.array]:
     """Forecasts on a list of time series.
 
     Args:
@@ -443,92 +272,7 @@ class TimesFm:
     Raises:
     ValueError: If the checkpoint is not properly loaded.
     """
-    if not self._train_state or not self._model:
-      raise ValueError(
-          "Checkpoint not loaded. Call `load_from_checkpoint` before"
-          " `forecast`.")
-    if forecast_context_len is None:
-      forecast_context_len = self.context_len
-    inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
-    inp_min = np.min([np.min(ts) for ts in inputs])
-
-    if window_size is not None:
-      new_inputs = []
-      for ts in inputs:
-        new_inputs.extend(moving_average(ts, window_size))
-      inputs = new_inputs
-
-    if freq is None:
-      logging.info("No frequency provided via `freq`. Default to high (0).")
-      freq = [0] * len(inputs)
-
-    input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
-    with base_layer.JaxContext.new_context(hparams=self._eval_context):
-      mean_outputs = []
-      full_outputs = []
-      assert input_ts.shape[0] % self.global_batch_size == 0
-      for i in range(input_ts.shape[0] // self.global_batch_size):
-        input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
-                                         self.global_batch_size])
-        input_padding_in = jnp.array(
-            input_padding[i * self.global_batch_size:(i + 1) *
-                          self.global_batch_size],)
-        inp_freq_in = jnp.array(
-            inp_freq[i * self.global_batch_size:(i + 1) *
-                     self.global_batch_size, :],
-            dtype=jnp.int32,
-        )
-        pmapped_inputs = NestedMap({
-            "input_ts":
-                es.jax_einshape(
-                    "(db)...->db...",
-                    input_ts_in,
-                    d=self.num_devices,
-                ),
-            "input_padding":
-                es.jax_einshape(
-                    "(db)...->db...",
-                    input_padding_in,
-                    d=self.num_devices,
-                ),
-            "date_features":
-                None,
-            "freq":
-                es.jax_einshape(
-                    "(db)...->db...",
-                    inp_freq_in,
-                    d=self.num_devices,
-                ),
-        })
-        mean_output, full_output = self._pmapped_decode(pmapped_inputs)
-        if not return_forecast_on_context:
-          mean_output = mean_output[:, :, self._horizon_start:, ...]
-          full_output = full_output[:, :, self._horizon_start:, ...]
-        mean_output = es.jax_einshape("db...->(db)...",
-                                      mean_output,
-                                      d=self.num_devices)
-        full_output = es.jax_einshape("db...->(db)...",
-                                      full_output,
-                                      d=self.num_devices)
-        mean_output = np.array(mean_output)
-        full_output = np.array(full_output)
-        mean_outputs.append(mean_output)
-        full_outputs.append(full_output)
-
-    mean_outputs = np.concatenate(mean_outputs, axis=0)
-    full_outputs = np.concatenate(full_outputs, axis=0)
-
-    if pmap_pad > 0:
-      mean_outputs = mean_outputs[:-pmap_pad, ...]
-      full_outputs = full_outputs[:-pmap_pad, ...]
-
-    if window_size is not None:
-      mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
-      full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
-    if inp_min >= 0 and truncate_negative:
-      mean_outputs = np.maximum(mean_outputs, 0.0)
-      full_outputs = np.maximum(full_outputs, 0.0)
-    return mean_outputs, full_outputs
+    raise NotImplementedError("`forecast` is not implemented.")
 
   def forecast_with_covariates(
       self,
@@ -663,7 +407,7 @@ class TimesFm:
       ]
       per_instance_stats = None
       if normalize_xreg_target_per_input:
-        targets, per_instance_stats = _normalize(targets)
+        targets, per_instance_stats = normalize(targets)
       xregs = xreg_lib.BatchedInContextXRegLinear(
           targets=targets,
           train_lens=train_lens,
@@ -686,7 +430,7 @@ class TimesFm:
           assert_covariate_shapes=True,
       )
       if normalize_xreg_target_per_input:
-        xregs = _renormalize(xregs, per_instance_stats)
+        xregs = renormalize(xregs, per_instance_stats)
       outputs = [
           (mean_output[self._horizon_start:(self._horizon_start + test_len)] +
            xreg)
@@ -701,7 +445,7 @@ class TimesFm:
       ]
       per_instance_stats = None
       if normalize_xreg_target_per_input:
-        targets, per_instance_stats = _normalize(targets)
+        targets, per_instance_stats = normalize(targets)
       xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear(
           targets=targets,
           train_lens=train_lens,
@@ -739,7 +483,7 @@ class TimesFm:
           for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
       ]
       if normalize_xreg_target_per_input:
-        outputs = _renormalize(outputs, per_instance_stats)
+        outputs = renormalize(outputs, per_instance_stats)
 
     return outputs, xregs
 
@@ -825,12 +569,11 @@ class TimesFm:
     )
     fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
 
-    if self._model.quantiles is not None:
-      for i, q in enumerate(self._model.quantiles):
-        q_col = f"{model_name}-q-{q}"
-        fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
-                                       1 + i].reshape(-1, 1)
-        if q == 0.5:
-          fcst_df[model_name] = fcst_df[q_col]
+    for i, q in enumerate(self.quantiles):
+      q_col = f"{model_name}-q-{q}"
+      fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
+                                     1 + i].reshape(-1, 1)
+      if q == 0.5:
+        fcst_df[model_name] = fcst_df[q_col]
     logging.info("Finished creating output dataframe.")
     return fcst_df

src/timesfm/timesfm_jax.py

@@ -0,0 +1,358 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TimesFM JAX forecast API for inference."""
+
+import logging
+import multiprocessing
+import time
+from os import path
+from typing import Any, Sequence
+
+import einshape as es
+import jax
+import jax.numpy as jnp
+import numpy as np
+from huggingface_hub import snapshot_download
+
+from paxml import checkpoints, tasks_lib
+from praxis import base_hyperparams, base_layer, pax_fiddle, py_utils, pytypes
+from praxis.layers import normalizations, transformers
+from timesfm import timesfm_base
+from timesfm import patched_decoder
+
+instantiate = base_hyperparams.instantiate
+NestedMap = py_utils.NestedMap
+JTensor = pytypes.JTensor
+
+_TOL = 1e-6
+
+
+class TimesFmJax(timesfm_base.TimesFmBase):
+  """TimesFM forecast API for inference.
+
+  This class is the scaffolding for calling TimesFM forecast. To properly use:
+    1. Create an instance with the correct hyperparameters of a TimesFM model.
+    2. Call `load_from_checkpoint` to load a compatible checkpoint.
+    3. Call `forecast` for inference.
+
+  Given the model size, this API does not shard the model weights for SPMD. All
+  parallelism happens on the data dimension.
+
+  Compilation happens during the first time `forecast` is called and uses the
+  `per_core_batch_size` to set and freeze the input signature. Subsequent calls
+  to `forecast` reflect the actual inference latency.
+  """
+
+  def _get_sample_inputs(self):
+    return {
+        "input_ts":
+            jnp.zeros(
+                (
+                    self.per_core_batch_size,
+                    self.context_len + self.output_patch_len,
+                ),
+                dtype=jnp.float32,
+            ),
+        "input_padding":
+            jnp.zeros(
+                (
+                    self.per_core_batch_size,
+                    self.context_len + self.output_patch_len,
+                ),
+                dtype=jnp.float32,
+            ),
+        "freq":
+            jnp.zeros(
+                (
+                    self.per_core_batch_size,
+                    1,
+                ),
+                dtype=jnp.int32,
+            ),
+    }
+
+  def __post_init__(self):
+    self.num_cores = jax.local_device_count(self.backend)
+    self.global_batch_size = self.per_core_batch_size * self.num_cores
+    self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
+    self._pmapped_decode = None
+    self._model = None
+    self._train_state = None
+
+  def load_from_checkpoint(
+      self,
+      checkpoint: timesfm_base.TimesFmCheckpoint,
+  ) -> None:
+    """Loads a checkpoint and compiles the decoder."""
+    checkpoint_type = (checkpoints.CheckpointType.FLAX
+                       if checkpoint.type is None else checkpoint.type)
+    checkpoint_path = checkpoint.path
+    step = checkpoint.step
+    repo_id = checkpoint.huggingface_repo_id
+    if checkpoint_path is None:
+      checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints")
+    # Rewrite the devices for Jax.
+    self.mesh_shape = [1, self.num_cores, 1]
+    self.mesh_name = ["replica", "data", "mdl"]
+
+    self.model_p = pax_fiddle.Config(
+        patched_decoder.PatchedTimeSeriesDecoder,
+        name="patched_decoder",
+        horizon_len=self.output_patch_len,
+        patch_len=self.input_patch_len,
+        model_dims=self.model_dims,
+        hidden_dims=self.model_dims,
+        residual_block_tpl=pax_fiddle.Config(patched_decoder.ResidualBlock),
+        quantiles=self.quantiles,
+        use_freq=True,
+        stacked_transformer_params_tpl=pax_fiddle.Config(
+            transformers.StackedTransformer,
+            num_heads=self.num_heads,
+            num_layers=self.num_layers,
+            transformer_layer_params_tpl=pax_fiddle.Config(
+                transformers.Transformer,
+                ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
+            ),
+        ),
+    )
+
+    self._key1, self._key2 = jax.random.split(jax.random.PRNGKey(42))
+    self._model = None
+    self._train_state = None
+    self._pmapped_decode = None
+    self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
+    try:
+      multiprocessing.set_start_method("spawn")
+    except RuntimeError:
+      print("Multiprocessing context has already been set.")
+    # Download the checkpoint from Hugging Face Hub if not given
+
+    #  Initialize the model weights.
+    self._logging("Constructing model weights.")
+    start_time = time.time()
+    self._model = instantiate(self.model_p)
+    var_weight_hparams = self._model.abstract_init_with_metadata(
+        self._get_sample_inputs(), do_eval=True)
+    train_state_partition_specs = tasks_lib.create_state_partition_specs(
+        var_weight_hparams,
+        mesh_shape=self.mesh_shape,
+        mesh_axis_names=self.mesh_name,
+        discard_opt_states=True,
+        learners=None,
+    )
+    train_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
+        var_weight_hparams,
+        discard_opt_states=True,
+        learners=None,
+    )
+    self._logging(
+        f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
+
+    # Load the model weights.
+    self._logging(f"Restoring checkpoint from {checkpoint_path}.")
+    start_time = time.time()
+    self._train_state = checkpoints.restore_checkpoint(
+        train_state_local_shapes,
+        checkpoint_dir=checkpoint_path,
+        checkpoint_type=checkpoint_type,
+        state_specs=train_state_partition_specs,
+        step=step,
+    )
+    self._logging(
+        f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
+    self.jit_decode()
+
+  def jit_decode(self):
+    """Jitting decoding function."""
+
+    # Initialize and jit the decode fn.
+    def _decode(inputs):
+      assert self._model is not None
+      assert self._train_state is not None
+      return self._model.apply(
+          self._train_state.mdl_vars,
+          inputs,
+          horizon_len=self.horizon_len,
+          output_patch_len=self.output_patch_len,
+          max_len=self.context_len,
+          return_forecast_on_context=True,
+          rngs={
+              base_layer.PARAMS: self._key1,
+              base_layer.RANDOM: self._key2,
+          },
+          method=self._model.decode,
+      )
+
+    self._logging("Jitting decoding.")
+    start_time = time.time()
+    self._pmapped_decode = jax.pmap(
+        _decode,
+        axis_name="batch",
+        devices=jax.devices(self.backend),
+        backend=self.backend,
+        axis_size=self.num_cores,
+    )
+    with base_layer.JaxContext.new_context(hparams=self._eval_context):
+      _ = self._pmapped_decode(
+          NestedMap({
+              "input_ts":
+                  jnp.zeros(
+                      (
+                          self.num_cores,
+                          self.per_core_batch_size,
+                          self.context_len,
+                      ),
+                      dtype=jnp.float32,
+                  ),
+              "input_padding":
+                  jnp.zeros(
+                      (
+                          self.num_cores,
+                          self.per_core_batch_size,
+                          self.context_len + self.horizon_len,
+                      ),
+                      dtype=jnp.float32,
+                  ),
+              "date_features":
+                  None,
+              "freq":
+                  jnp.zeros(
+                      (self.num_cores, self.per_core_batch_size, 1),
+                      dtype=jnp.int32,
+                  ),
+          }))
+    self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
+
+  def forecast(
+      self,
+      inputs: Sequence[Any],
+      freq: Sequence[int] | None = None,
+      window_size: int | None = None,
+      forecast_context_len: int | None = None,
+      return_forecast_on_context: bool = False,
+      truncate_negative: bool = False,
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Forecasts on a list of time series.
+
+    Args:
+      inputs: list of time series forecast contexts. Each context time series
+        should be in a format convertible to JTensor by `jnp.array`.
+      freq: frequency of each context time series. 0 for high frequency
+        (default), 1 for medium, and 2 for low. Notice this is different from
+        the `freq` required by `forecast_on_df`.
+      window_size: window size of trend + residual decomposition. If None then
+        we do not do decomposition.
+      forecast_context_len: optional max context length.
+      return_forecast_on_context: True to return the forecast on the context
+        when available, i.e. after the first input patch.
+      truncate_negative: truncate to only non-negative values if all the contexts
+        have non-negative values.
+
+    Returns:
+    A tuple for JTensors:
+    - the mean forecast of size (# inputs, # forecast horizon),
+    - the full forecast (mean + quantiles) of size
+        (# inputs,  # forecast horizon, 1 + # quantiles).
+
+    Raises:
+    ValueError: If the checkpoint is not properly loaded.
+    """
+    if not self._train_state or not self._model:
+      raise ValueError(
+          "Checkpoint not loaded. Call `load_from_checkpoint` before"
+          " `forecast`.")
+    if forecast_context_len is None:
+      fcontext_len = self.context_len
+    else:
+      fcontext_len = forecast_context_len
+    inputs = [np.array(ts)[-fcontext_len:] for ts in inputs]
+    inp_min = np.min([np.min(ts) for ts in inputs])
+
+    if window_size is not None:
+      new_inputs = []
+      for ts in inputs:
+        new_inputs.extend(timesfm_base.moving_average(ts, window_size))
+      inputs = new_inputs
+
+    if freq is None:
+      logging.info("No frequency provided via `freq`. Default to high (0).")
+      freq = [0] * len(inputs)
+
+    input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
+    with base_layer.JaxContext.new_context(hparams=self._eval_context):
+      mean_outputs = []
+      full_outputs = []
+      assert input_ts.shape[0] % self.global_batch_size == 0
+      for i in range(input_ts.shape[0] // self.global_batch_size):
+        input_ts_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
+                                         self.global_batch_size])
+        input_padding_in = jnp.array(
+            input_padding[i * self.global_batch_size:(i + 1) *
+                          self.global_batch_size],)
+        inp_freq_in = jnp.array(
+            inp_freq[i * self.global_batch_size:(i + 1) *
+                     self.global_batch_size, :],
+            dtype=jnp.int32,
+        )
+        pmapped_inputs = NestedMap({
+            "input_ts":
+                es.jax_einshape(
+                    "(db)...->db...",
+                    input_ts_in,
+                    d=self.num_cores,
+                ),
+            "input_padding":
+                es.jax_einshape(
+                    "(db)...->db...",
+                    input_padding_in,
+                    d=self.num_cores,
+                ),
+            "date_features":
+                None,
+            "freq":
+                es.jax_einshape(
+                    "(db)...->db...",
+                    inp_freq_in,
+                    d=self.num_cores,
+                ),
+        })
+        mean_output, full_output = self._pmapped_decode(pmapped_inputs)
+        if not return_forecast_on_context:
+          mean_output = mean_output[:, :, self._horizon_start:, ...]
+          full_output = full_output[:, :, self._horizon_start:, ...]
+        mean_output = es.jax_einshape("db...->(db)...",
+                                      mean_output,
+                                      d=self.num_cores)
+        full_output = es.jax_einshape("db...->(db)...",
+                                      full_output,
+                                      d=self.num_cores)
+        mean_output = np.array(mean_output)
+        full_output = np.array(full_output)
+        mean_outputs.append(mean_output)
+        full_outputs.append(full_output)
+
+    mean_outputs = np.concatenate(mean_outputs, axis=0)
+    full_outputs = np.concatenate(full_outputs, axis=0)
+
+    if pmap_pad > 0:
+      mean_outputs = mean_outputs[:-pmap_pad, ...]
+      full_outputs = full_outputs[:-pmap_pad, ...]
+
+    if window_size is not None:
+      mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
+      full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
+    if inp_min >= 0 and truncate_negative:
+      mean_outputs = np.maximum(mean_outputs, 0.0)
+      full_outputs = np.maximum(full_outputs, 0.0)
+    return mean_outputs, full_outputs

src/timesfm/timesfm_torch.py

@@ -0,0 +1,171 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TimesFM pytorch forecast API for inference."""
+
+import logging
+from os import path
+from typing import Any, Sequence
+
+import numpy as np
+import torch
+from huggingface_hub import snapshot_download
+from timesfm import timesfm_base
+
+from . import pytorch_patched_decoder as ppd
+
+_TOL = 1e-6
+
+
+class TimesFmTorch(timesfm_base.TimesFmBase):
+  """TimesFM forecast API for inference."""
+
+  def __post_init__(self):
+    self._model_config = ppd.TimesFMConfig(
+        num_layers=self.num_layers,
+        num_heads=self.num_heads,
+        hidden_size=self.model_dims,
+        intermediate_size=self.model_dims,
+        patch_len=self.input_patch_len,
+        horizon_len=self.output_patch_len,
+        head_dim=self.model_dims // self.num_heads,
+        quantiles=self.quantiles,
+    )
+    self._model = None
+    self.num_cores = 1
+    self.global_batch_size = self.per_core_batch_size
+    self._device = torch.device("cuda:0" if (
+        torch.cuda.is_available() and self.backend == "gpu") else "cpu")
+
+  def load_from_checkpoint(
+      self,
+      checkpoint: timesfm_base.TimesFmCheckpoint,
+  ) -> None:
+    """Loads a checkpoint and compiles the decoder."""
+    checkpoint_path = checkpoint.path
+    repo_id = checkpoint.huggingface_repo_id
+    if checkpoint_path is None:
+      checkpoint_path = path.join(snapshot_download(repo_id),
+                                  "torch_model.ckpt")
+    self._model = ppd.PatchedTimeSeriesDecoder(self._model_config)
+    loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+    logging.info("Loading checkpoint from %s", checkpoint_path)
+    self._model.load_state_dict(loaded_checkpoint)
+    logging.info("Sending checkpoint to device %s", f"{self._device}")
+    self._model.to(self._device)
+    self._model.eval()
+    # TODO: add compilation.
+
+  def forecast(
+      self,
+      inputs: Sequence[Any],
+      freq: Sequence[int] | None = None,
+      window_size: int | None = None,
+      forecast_context_len: int | None = None,
+      return_forecast_on_context: bool = False,
+      truncate_negative: bool = False,
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Forecasts on a list of time series.
+
+        Args:
+          inputs: list of time series forecast contexts. Each context time series
+            should be in a format convertible to JTensor by `jnp.array`.
+          freq: frequency of each context time series. 0 for high frequency
+            (default), 1 for medium, and 2 for low. Notice this is different from
+            the `freq` required by `forecast_on_df`.
+          window_size: window size of trend + residual decomposition. If None then
+            we do not do decomposition.
+          forecast_context_len: optional max context length.
+          return_forecast_on_context: True to return the forecast on the context
+            when available, i.e. after the first input patch.
+          truncate_negative: truncate to only non-negative values if all the contexts
+            have non-negative values.
+
+        Returns:
+        A tuple for JTensors:
+        - the mean forecast of size (# inputs, # forecast horizon),
+        - the full forecast (mean + quantiles) of size
+            (# inputs,  # forecast horizon, 1 + # quantiles).
+
+        Raises:
+        ValueError: If the checkpoint is not properly loaded.
+        """
+    if not self._model:
+      raise ValueError(
+          "Checkpoint not loaded. Call `load_from_checkpoint` before"
+          " `forecast`.")
+    if forecast_context_len is None:
+      fcontext_len = self.context_len
+    else:
+      fcontext_len = forecast_context_len
+    inputs = [np.array(ts)[-fcontext_len:] for ts in inputs]
+    inp_min = np.min([np.min(ts) for ts in inputs])
+
+    if window_size is not None:
+      new_inputs = []
+      for ts in inputs:
+        new_inputs.extend(timesfm_base.moving_average(ts, window_size))
+      inputs = new_inputs
+
+    if freq is None:
+      logging.info("No frequency provided via `freq`. Default to high (0).")
+      freq = [0] * len(inputs)
+
+    input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
+    with torch.no_grad():
+      mean_outputs = []
+      full_outputs = []
+      assert input_ts.shape[0] % self.global_batch_size == 0
+      for i in range(input_ts.shape[0] // self.global_batch_size):
+        input_ts_in = torch.from_numpy(
+            np.array(input_ts[i * self.global_batch_size:(i + 1) *
+                              self.global_batch_size],
+                     dtype=np.float32)).to(self._device)
+        input_padding_in = torch.from_numpy(
+            np.array(input_padding[i * self.global_batch_size:(i + 1) *
+                                   self.global_batch_size],
+                     dtype=np.float32)).to(self._device)
+        inp_freq_in = torch.from_numpy(
+            np.array(inp_freq[
+                i * self.global_batch_size:(i + 1) * self.global_batch_size,
+                :,
+            ],
+                     dtype=np.int32)).long().to(self._device)
+        mean_output, full_output = self._model.decode(
+            input_ts=input_ts_in,
+            paddings=input_padding_in,
+            freq=inp_freq_in,
+            horizon_len=self.horizon_len,
+            return_forecast_on_context=return_forecast_on_context,
+        )
+        mean_output = mean_output.detach().cpu().numpy()
+        full_output = full_output.detach().cpu().numpy()
+        mean_output = np.array(mean_output)
+        full_output = np.array(full_output)
+        mean_outputs.append(mean_output)
+        full_outputs.append(full_output)
+
+    mean_outputs = np.concatenate(mean_outputs, axis=0)
+    full_outputs = np.concatenate(full_outputs, axis=0)
+
+    if pmap_pad > 0:
+      mean_outputs = mean_outputs[:-pmap_pad, ...]
+      full_outputs = full_outputs[:-pmap_pad, ...]
+
+    if window_size is not None:
+      mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
+      full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
+    if inp_min >= 0 and truncate_negative:
+      mean_outputs = np.maximum(mean_outputs, 0.0)
+      full_outputs = np.maximum(full_outputs, 0.0)
+    return mean_outputs, full_outputs

src/timesfm_torch/README.md

@@ -1,9 +0,0 @@
-# Official Pytorch implementation of TimesFM
-
-TimesFM  (Time Series Foundation Model) is a pretrained time-series foundation model developed by Google
-Research for time-series forecasting.
-
-* Paper: [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688), to appear in ICML 2024.
-* [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/)
-
-## Stay tuned for all of the functionalities as that of the pax version.