Merge pull request #91 from google-research/test_650786820

In context XReg.

environment.yml

@@ -15,3 +15,4 @@ dependencies:
     - paxml
     - jax[cuda12]==0.4.26
     - einshape
+    - scikit-learn

environment_cpu.yml

@@ -15,3 +15,4 @@ dependencies:
     - paxml
     - jax[cpu]==0.4.26
     - einshape
+    - scikit-learn

experiments/baselines/timegpt_pipeline.py

@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional, Tuple
 import os
-import pandas as pd
+from time import time
+from typing import List, Optional, Tuple
+from dotenv import load_dotenv
 from gluonts.time_feature.seasonality import get_seasonality as _get_seasonality
+from nixtla import NixtlaClient
+import pandas as pd
 from tqdm import tqdm
 from utilsforecast.processing import (
     backtest_splits,
@@ -25,17 +28,15 @@ from utilsforecast.processing import (
     take_rows,
     vertical_concat,
 )
-from time import time
-from dotenv import load_dotenv
-from nixtla import NixtlaClient
 
 
 def get_seasonality(freq: str) -> int:
   return _get_seasonality(freq, seasonalities={"D": 7})
 
 
-def maybe_convert_col_to_datetime(df: pd.DataFrame,
+def maybe_convert_col_to_datetime(
-                                  col_name: str) -> pd.DataFrame:
+    df: pd.DataFrame, col_name: str
+) -> pd.DataFrame:
   if not pd.api.types.is_datetime64_any_dtype(df[col_name]):
     df = df.copy()
     df[col_name] = pd.to_datetime(df[col_name])
@@ -63,15 +64,14 @@ def zero_pad_time_series(df, freq, min_length=36):
           end=start_date,
           periods=min_length - len(subset) + 1,
           freq=freq,  # 'MS' for month start
-      )[:-1]  # Exclude the start_date itself
+      )[
+          :-1
+      ]  # Exclude the start_date itself
 
       # 2c. Create padding data
-      padding_df = pd.DataFrame({
+      padding_df = pd.DataFrame(
-          "ds": padding_dates,
+          {"ds": padding_dates, "unique_id": unique_id, "y": 0}  # Zero padding
-          "unique_id": unique_id,
+      )
-          "y": 0
-      }  # Zero padding
-                               )
 
       # 2d. Combine original and padding data, and append to the list
       padded_data.append(pd.concat([padding_df, subset]).sort_values("ds"))
@@ -83,8 +83,9 @@ def zero_pad_time_series(df, freq, min_length=36):
 
 class Forecaster:
   """Borrowed from
-    https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
+
-    """
+  https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
+  """
 
   def forecast(
       self,
@@ -120,7 +121,8 @@ class Forecaster:
     for _, (cutoffs, train, valid) in tqdm(enumerate(splits)):
       if len(valid.columns) > 3:
         raise NotImplementedError(
-            "Cross validation with exogenous variables is not yet supported.")
+            "Cross validation with exogenous variables is not yet supported."
+        )
       y_pred = self.forecast(
           df=train,
           h=h,
@@ -134,9 +136,10 @@ class Forecaster:
       )
       if result.shape[0] < valid.shape[0]:
         raise ValueError(
-            "Cross validation result produced less results than expected. "
+            "Cross validation result produced less results than expected."
-            "Please verify that the frequency parameter (freq) matches your series' "
+            " Please verify that the frequency parameter (freq) matches your"
-            "and that there aren't any missing periods.")
+            " series' and that there aren't any missing periods."
+        )
       results.append(result)
     out = vertical_concat(results)
     out = drop_index_if_pandas(out)
@@ -148,9 +151,10 @@ class Forecaster:
 
 class TimeGPT(Forecaster):
   """Borrowed from
-    https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
+
-    We modify the class to take care of edge cases.
+  https://github.com/Nixtla/nixtla/tree/main/experiments/foundation-time-series-arena/xiuhmolpilli/models.
-    """
+  We modify the class to take care of edge cases.
+  """
 
   def __init__(
       self,
@@ -199,7 +203,7 @@ class TimeGPT(Forecaster):
       all_unique_ids = df["unique_id"].unique()
       all_fcst_df = []
       for i in range(0, len(all_unique_ids), chunk_size):
-        chunk_ids = all_unique_ids[i:i + chunk_size]
+        chunk_ids = all_unique_ids[i : i + chunk_size]
         chunk_df = df[df["unique_id"].isin(chunk_ids)]
         fct_chunk_df = client.forecast(
             df=chunk_df,
@@ -236,11 +240,13 @@ def run_timegpt(
     chunk_size = 5000
   else:
     chunk_size = None
-  fcsts_df = model.forecast(df=padded_train_df,
+  fcsts_df = model.forecast(
-                            h=horizon,
+      df=padded_train_df,
-                            level=level,
+      h=horizon,
-                            freq=freq,
+      level=level,
-                            chunk_size=chunk_size)
+      freq=freq,
+      chunk_size=chunk_size,
+  )
   total_time = time() - init_time
   # In case levels are not returned we replace the levels with the mean predictions.
   # Note that this does not affect the results table as we only compare on point

experiments/environment.yml

@@ -25,3 +25,4 @@ dependencies:
     - python-dotenv
     - nixtla>=0.5.1
     - rich
+    - scikit-learn

experiments/environment_cpu.yml

@@ -25,3 +25,4 @@ dependencies:
     - python-dotenv
     - nixtla>=0.5.1
     - rich
+    - scikit-learn

experiments/extended_benchmarks/run_timegpt.py

@@ -11,6 +11,7 @@
 # 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 timegpt."""
 
 import os
@@ -20,10 +21,11 @@ import time
 from absl import flags
 import numpy as np
 import pandas as pd
-from ..baselines.timegpt_pipeline import run_timegpt
 
+from ..baselines.timegpt_pipeline import run_timegpt
 from .utils import ExperimentHandler
 
+
 dataset_names = [
     "m1_monthly",
     "m1_quarterly",
@@ -61,6 +63,7 @@ _MODEL_NAME = flags.DEFINE_string(
 )
 _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
 
+
 QUANTILES = list(np.arange(1, 10) / 10.0)
 
 
@@ -87,9 +90,9 @@ def main():
     )
     time_df = pd.DataFrame({"time": [total_time], "model": model_name})
     fcsts_df = exp.fcst_from_level_to_quantiles(fcsts_df, model_name)
-    results = exp.evaluate_from_predictions(models=[model_name],
+    results = exp.evaluate_from_predictions(
-                                            fcsts_df=fcsts_df,
+        models=[model_name], fcsts_df=fcsts_df, times_df=time_df
-                                            times_df=time_df)
+    )
     print(results, flush=True)
     results_list.append(results)
     results_full = pd.concat(results_list)

experiments/extended_benchmarks/run_timesfm.py

@@ -11,6 +11,7 @@
 # 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
@@ -25,6 +26,7 @@ import timesfm
 
 from .utils import ExperimentHandler
 
+
 dataset_names = [
     "m1_monthly",
     "m1_quarterly",
@@ -72,14 +74,16 @@ context_dict = {
     "m4_yearly": 64,
 }
 
-_MODEL_PATH = flags.DEFINE_string("model_path", "/home/timesfm_q10_20240501",
+_MODEL_PATH = flags.DEFINE_string(
-                                  "Path to model")
+    "model_path", "/home/timesfm_q10_20240501", "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)
 
 
@@ -123,9 +127,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],
+    results = exp.evaluate_from_predictions(
-                                            fcsts_df=fcsts_df,
+        models=[model_name], fcsts_df=fcsts_df, times_df=time_df
-                                            times_df=time_df)
+    )
     print(results, flush=True)
     results_list.append(results)
     results_full = pd.concat(results_list)

experiments/extended_benchmarks/utils.py

@@ -11,6 +11,7 @@
 # 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.
+
 """Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py."""
 
 from functools import partial
@@ -45,9 +46,11 @@ def quantile_loss(
     target_col: str = "y",
 ) -> pd.DataFrame:
   delta_y = df[models].sub(df[target_col], axis=0)
-  res = (np.maximum(q * delta_y,
+  res = (
-                    (q - 1) * delta_y).groupby(df[id_col],
+      np.maximum(q * delta_y, (q - 1) * delta_y)
-                                               observed=True).mean())
+      .groupby(df[id_col], observed=True)
+      .mean()
+  )
   res.index.name = id_col
   res = res.reset_index()
   return res
@@ -63,8 +66,10 @@ class ExperimentHandler:
       models_dir: str = "./models",
   ):
     if dataset not in gluonts_datasets:
-      raise Exception(f"dataset {dataset} not found in gluonts "
+      raise Exception(
-                      f"available datasets: {', '.join(gluonts_datasets)}")
+          f"dataset {dataset} not found in gluonts "
+          f"available datasets: {', '.join(gluonts_datasets)}"
+      )
     self.dataset = dataset
     self.quantiles = quantiles
     self.level = self._transform_quantiles_to_levels(quantiles)
@@ -75,8 +80,10 @@ class ExperimentHandler:
     gluonts_dataset = get_dataset(self.dataset)
     self.horizon = gluonts_dataset.metadata.prediction_length
     if self.horizon is None:
-      raise Exception(f"horizon not found for dataset {self.dataset} "
+      raise Exception(
-                      "experiment cannot be run")
+          f"horizon not found for dataset {self.dataset} "
+          "experiment cannot be run"
+      )
     self.freq = gluonts_dataset.metadata.freq
     # get_seasonality() returns 1 for freq='D', override this to 7. This significantly improves the accuracy of
     # statistical models on datasets like m5/nn5_daily. The models like AutoARIMA/AutoETS can still set
@@ -115,8 +122,9 @@ class ExperimentHandler:
 
   @staticmethod
   def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]:
-    level = [int(100 - 200 * q) for q in quantiles if q < 0.5
+    level = [
-            ]  # in this case mean=mediain
+        int(100 - 200 * q) for q in quantiles if q < 0.5
+    ]  # in this case mean=mediain
     level = sorted(list(set(level)))
     return level
 
@@ -145,8 +153,9 @@ class ExperimentHandler:
       last_n: int | None = None,
   ) -> pd.DataFrame:
     with multiprocessing.Pool(os.cpu_count()) as pool:  # Create a process pool
-      results = pool.map(parallel_transform, zip(gluonts_dataset,
+      results = pool.map(
-                                                 repeat(last_n)))
+          parallel_transform, zip(gluonts_dataset, repeat(last_n))
+      )
     df = pd.concat(results)
     df = df.reset_index(drop=True)
     return df
@@ -168,8 +177,9 @@ class ExperimentHandler:
   def save_dataframe(self, df: pd.DataFrame, file_name: str):
     df.to_csv(f"{self.results_dir}/{file_name}", index=False)
 
-  def save_results(self, fcst_df: pd.DataFrame, total_time: float,
+  def save_results(
-                   model_name: str):
+      self, fcst_df: pd.DataFrame, total_time: float, model_name: str
+  ):
     self.save_dataframe(
         fcst_df,
         f"{model_name}-{self.dataset}-fcst.csv",
@@ -205,21 +215,23 @@ class ExperimentHandler:
     times_df = []
     for model in models:
       fcst_method_df = pd.read_csv(
-          f"{self.results_dir}/{model}-{self.dataset}-fcst.csv").set_index(
+          f"{self.results_dir}/{model}-{self.dataset}-fcst.csv"
-              ["unique_id", "ds"])
+      ).set_index(["unique_id", "ds"])
       fcsts_df.append(fcst_method_df)
       time_method_df = pd.read_csv(
-          f"{self.results_dir}/{model}-{self.dataset}-time.csv")
+          f"{self.results_dir}/{model}-{self.dataset}-time.csv"
+      )
       times_df.append(time_method_df)
     fcsts_df = pd.concat(fcsts_df, axis=1).reset_index()
     fcsts_df["ds"] = pd.to_datetime(fcsts_df["ds"])
     times_df = pd.concat(times_df)
-    return self.evaluate_from_predictions(models=models,
+    return self.evaluate_from_predictions(
-                                          fcsts_df=fcsts_df,
+        models=models, fcsts_df=fcsts_df, times_df=times_df
-                                          times_df=times_df)
+    )
 
-  def evaluate_from_predictions(self, models: List[str], fcsts_df: pd.DataFrame,
+  def evaluate_from_predictions(
-                                times_df: pd.DataFrame) -> pd.DataFrame:
+      self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame
+  ) -> pd.DataFrame:
     test_df = self.test_df
     train_df = self.train_df
     test_df = test_df.merge(fcsts_df, how="left")
@@ -250,9 +262,9 @@ class ExperimentHandler:
     eval_prob_df["metric"] = "scaled_crps"
     eval_df = pd.concat([eval_df, eval_prob_df]).reset_index(drop=True)
     eval_df = eval_df.groupby("metric").mean(numeric_only=True).reset_index()
-    eval_df = eval_df.melt(id_vars="metric",
+    eval_df = eval_df.melt(
-                           value_name="value",
+        id_vars="metric", value_name="value", var_name="model"
-                           var_name="model")
+    )
     times_df.insert(0, "metric", "time")
     times_df = times_df.rename(columns={"time": "value"})
     eval_df = pd.concat([eval_df, times_df])

experiments/long_horizon_benchmarks/run_eval.py

@@ -11,6 +11,7 @@
 # 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
@@ -23,32 +24,42 @@ import numpy as np
 import pandas as pd
 from paxml import checkpoints
 import timesfm
+from timesfm import data_loader
 import torch
 import tqdm
-from timesfm import data_loader
+
 
 FLAGS = flags.FLAGS
 
-_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64,
+_BATCH_SIZE = flags.DEFINE_integer(
-                                   "Batch size for the randomly sampled batch")
+    "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",
+_MODEL_PATH = flags.DEFINE_string(
-                                  "The name of the dataset.")
+    "model_path", "./timesfm_q10_20240501", "The name of the dataset."
-_DATETIME_COL = flags.DEFINE_string("datetime_col", "date",
+)
-                                    "Column having datetime.")
+_DATETIME_COL = flags.DEFINE_string(
-_NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None,
+    "datetime_col", "date", "Column having datetime."
-                                  "Column having numerical features.")
+)
-_CAT_COV_COLS = flags.DEFINE_list("cat_cov_cols", None,
+_NUM_COV_COLS = flags.DEFINE_list(
-                                  "Column having categorical features.")
+    "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 = flags.DEFINE_bool(
-                               "normalize data for eval or not")
+    "normalize", True, "normalize data for eval or not"
-_CONTEXT_LEN = flags.DEFINE_integer("context_len", 512,
+)
-                                    "Length of the context window")
+_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_DIR = flags.DEFINE_string(
-                                   "results directory")
+    "results_dir", "./results/long_horizon", "results directory"
+)
+
 
 DATA_DICT = {
     "ettm2": {
@@ -165,8 +176,9 @@ def eval():
       holiday=False,
       permute=False,
   )
-  eval_itr = dtl.tf_dataset(mode="test",
+  eval_itr = dtl.tf_dataset(
-                            shift=_PRED_LEN.value).as_numpy_iterator()
+      mode="test", shift=_PRED_LEN.value
+  ).as_numpy_iterator()
   model_path = _MODEL_PATH.value
   if model_path.startswith("amazon"):
     model = chronos.ChronosPipeline.from_pretrained(
@@ -201,9 +213,10 @@ def eval():
   for batch in tqdm.tqdm(eval_itr):
     past = batch[0]
     actuals = batch[3]
-    forecasts = get_forecasts(model_path, model, past, int_freq,
+    forecasts = get_forecasts(
-                              _PRED_LEN.value)
+        model_path, model, past, int_freq, _PRED_LEN.value
-    forecasts = forecasts[:, 0:actuals.shape[1]]
+    )
+    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

@@ -3,7 +3,7 @@
 [project]
 name = "timesfm"
 description = "Open weights time-series foundation model from Google Research."
-version = "0.0.1"
+version = "1.0.0"
 dependencies = [
     "einshape>=1.0.0",
     "paxml>=1.4.0",
@@ -11,6 +11,7 @@ dependencies = [
     "jax>=0.4.26",
     "numpy>=1.26.4",
     "pandas>=2.1.4",
+    "sklearn>=1.15.1",
 ]
 authors = [
      {name = "Rajat Sen", email = "senrajat@google.com"},

src/timesfm/__init__.py

@@ -11,6 +11,7 @@
 # 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

src/timesfm/data_loader.py

@@ -11,11 +11,13 @@
 # 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.
+
 """TF dataloaders for general timeseries datasets.
 
 The expected input format is csv file with a datetime index.
 """
 
+
 from absl import logging
 import numpy as np
 import pandas as pd
@@ -77,8 +79,9 @@ class TimeSeriesdata(object):
       self.data_df['ccol'] = np.zeros(self.data_df.shape[0])
       cat_cov_cols = ['ccol']
     self.data_df.fillna(0, inplace=True)
-    self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]),
+    self.data_df.set_index(
-                           inplace=True)
+        pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True
+    )
     self.num_cov_cols = num_cov_cols
     self.cat_cov_cols = cat_cov_cols
     self.ts_cols = ts_cols
@@ -91,16 +94,18 @@ class TimeSeriesdata(object):
             data_df_idx[-1] + pd.Timedelta(1, freq=freq),
             periods=pred_len + 1,
             freq=freq,
-        ))
+        )
+    )
     self.time_df = time_features.TimeCovariates(
-        date_index, holiday=holiday).get_covariates()
+        date_index, holiday=holiday
+    ).get_covariates()
     self.hist_len = hist_len
     self.pred_len = pred_len
     self.batch_size = batch_size
     self.freq = freq
     self.normalize = normalize
     self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose()
-    self.data_mat = self.data_mat[:, 0:self.test_range[1]]
+    self.data_mat = self.data_mat[:, 0 : self.test_range[1]]
     self.time_mat = self.time_df.to_numpy().transpose()
     self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose()
     self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols)
@@ -130,7 +135,7 @@ class TimeSeriesdata(object):
 
   def _normalize_data(self):
     self.scaler = StandardScaler()
-    train_mat = self.data_mat[:, self.train_range[0]:self.train_range[1]]
+    train_mat = self.data_mat[:, self.train_range[0] : self.train_range[1]]
     self.scaler = self.scaler.fit(train_mat.transpose())
     self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
 
@@ -248,8 +253,9 @@ class TimeSeriesdata(object):
       gen_fn = self.train_gen
     else:
       gen_fn = lambda: self.test_val_gen(mode, shift)
-    output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 +
+    output_types = tuple(
-                         [tf.int32] * 2)
+        [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2
+    )
     dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
     dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
     return dataset

src/timesfm/patched_decoder.py

@@ -11,6 +11,7 @@
 # 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.
@@ -35,6 +36,7 @@ from praxis.layers import normalizations
 from praxis.layers import stochastics
 from praxis.layers import transformers
 
+
 # PAX shortcuts
 NestedMap = py_utils.NestedMap
 JTensor = pytypes.JTensor
@@ -42,6 +44,7 @@ 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]
 
@@ -54,6 +57,7 @@ _FREQ = "freq"
 _OUTPUT_TOKENS = "output_tokens"
 _STATS = "stats"
 
+
 # Small numerical value.
 _TOLERANCE = 1e-7
 
@@ -84,16 +88,16 @@ def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor:
 class ResidualBlock(base_layer.BaseLayer):
   """Simple feedforward block with residual connection.
 
-    Attributes:
+  Attributes:
-      input_dims: input dimension.
+    input_dims: input dimension.
-      hidden_dims: hidden dimension.
+    hidden_dims: hidden dimension.
-      output_dims: output dimension.
+    output_dims: output dimension.
-      dropout_prob: dropout probability.
+    dropout_prob: dropout probability.
-      layer_norm: whether to use layer norm or not.
+    layer_norm: whether to use layer norm or not.
-      dropout_tpl: config for dropout.
+    dropout_tpl: config for dropout.
-      ln_tpl: config for layer norm.
+    ln_tpl: config for layer norm.
-      act_tpl: config for activation in hidden layer.
+    act_tpl: config for activation in hidden layer.
-    """
+  """
 
   input_dims: int = 0
   hidden_dims: int = 0
@@ -154,20 +158,21 @@ class ResidualBlock(base_layer.BaseLayer):
       return output + residual
 
 
-def _masked_mean_std(inputs: JTensor,
+def _masked_mean_std(
-                     padding: JTensor) -> Tuple[JTensor, JTensor]:
+    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.
+  It should exclude values where pad is 1.
 
-    Args:
+  Args:
-      inputs: A JAX array of shape [b, n, p].
+    inputs: A JAX array of shape [b, n, p].
-      padding: A JAX array of shape [b, n, p] with values 0 or 1.
+    padding: A JAX array of shape [b, n, p] with values 0 or 1.
 
-    Returns:
+  Returns:
-      A tuple containing the mean and standard deviation of arr. We return the
+    A tuple containing the mean and standard deviation of arr. We return the
-      statistics of the first patch with more than three non-padded values.
+    statistics of the first patch with more than three non-padded values.
-    """
+  """
   # Selecting the first pad with more than 3 unpadded values.
   pad_sum = jnp.sum(1 - padding, axis=2)
 
@@ -192,7 +197,7 @@ def _masked_mean_std(inputs: JTensor,
 
   # 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
@@ -211,22 +216,22 @@ def _create_quantiles() -> list[float]:
 class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
   """Patch decoder layer for time-series foundation model.
 
-    Attributes:
+  Attributes:
-      patch_len: length of input patches.
+    patch_len: length of input patches.
-      horizon_len: length of output patches. Referred to as `output_patch_len`
+    horizon_len: length of output patches. Referred to as `output_patch_len`
-        during inference.
+      during inference.
-      model_dims: model dimension of stacked transformer layer.
+    model_dims: model dimension of stacked transformer layer.
-      hidden_dims: hidden dimensions in fully connected layers.
+    hidden_dims: hidden dimensions in fully connected layers.
-      quantiles: list of quantiles for non prob model.
+    quantiles: list of quantiles for non prob model.
-      residual_block_tpl: config for residual block.
+    residual_block_tpl: config for residual block.
-      stacked_transformer_params_tpl: config for stacked transformer.
+    stacked_transformer_params_tpl: config for stacked transformer.
-      use_freq: whether to use frequency encoding.
+    use_freq: whether to use frequency encoding.
 
-    In all of what followed, except specified otherwise, B is batch size, T is
+  In all of what followed, except specified otherwise, B is batch size, T is
-    sequence length of time-series. N is the number of input patches that can be
+  sequence length of time-series. N is the number of input patches that can be
-    obtained from T. P is the input patch length and H is the horizon length. Q is
+  obtained from T. P is the input patch length and H is the horizon length. Q is
-    number of output logits. D is model dimension.
+  number of output logits. D is model dimension.
-    """
+  """
 
   patch_len: int = 0
   horizon_len: int = 0
@@ -235,7 +240,8 @@ 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:
@@ -270,8 +276,9 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
 
     self.create_child(
         "position_emb",
-        pax_fiddle.Config(layers.PositionalEmbedding,
+        pax_fiddle.Config(
-                          embedding_dims=self.model_dims),
+            layers.PositionalEmbedding, embedding_dims=self.model_dims
+        ),
     )
 
     if self.use_freq:
@@ -285,24 +292,27 @@ 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,
+      self, inputs: JTensor, patched_pads: JTensor
-      patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, 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,
+  def _reverse_transform(
-                         stats: Tuple[JTensor, JTensor]) -> JTensor:
+      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]
@@ -316,13 +326,19 @@ 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)
-    input_padding = jnp.where(
+    patched_pads = es.jax_einshape(
-        jnp.abs(input_ts - PAD_VAL) < _TOLERANCE, 1, input_padding)
+        "b(np)->bnp", input_padding, p=self.patch_len
-    patched_pads = es.jax_einshape("b(np)->bnp",
+    )
-                                   input_padding,
+    patched_inputs = jnp.where(
-                                   p=self.patch_len)
+        jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs
-    patched_inputs, stats = self._forward_transform(patched_inputs,
+    )
-                                                    patched_pads)
+    patched_pads = jnp.where(
+        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)
     concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1)
@@ -351,26 +367,25 @@ 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 = es.jax_einshape(
-                                output_ts,
+        "bn(hq)->bnhq", output_ts, q=num_outputs, h=self.horizon_len
-                                q=num_outputs,
+    )
-                                h=self.horizon_len)
     return self._reverse_transform(output_ts, stats)
 
   def __call__(self, inputs: NestedMap) -> NestedMap:
     """PatchTST call.
 
-        Args:
+    Args:
-          inputs: A NestedMap containing (1) input_ts: input sequence of shape [B,
+      inputs: A NestedMap containing (1) input_ts: input sequence of shape [B,
-            T] where T must be multiple of patch_length; (2) input_padding: that
+        T] where T must be multiple of patch_length; (2) input_padding: that
-            contains padding map.
+        contains padding map.
 
-        Returns:
+    Returns:
-          A nested map with two keys:
+      A nested map with two keys:
-          (1) 'output_tokens' of shape [B, N, D].
+      (1) 'output_tokens' of shape [B, N, D].
-          (2) 'output_ts' of shape [B, N, H, Q]
+      (2) 'output_ts' of shape [B, N, H, Q]
-          (3) 'stats' a Tuple of statistics for renormalization.
+      (3) 'stats' a Tuple of statistics for renormalization.
-        """
+    """
     input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING]
     num_outputs = len(self.quantiles) + 1
     model_input, patched_padding, stats, _ = self._preprocess_input(
@@ -385,11 +400,9 @@ 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({
+    return NestedMap(
-        _OUTPUT_TOKENS: model_output,
+        {_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats}
-        _OUTPUT_TS: output_ts,
+    )
-        _STATS: stats
-    })
 
   def decode(
       self,
@@ -397,25 +410,30 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
       horizon_len: int,
       output_patch_len: Optional[int] = None,
       max_len: int = 512,
+      return_forecast_on_context: bool = False,
   ) -> tuple[JTensor, JTensor]:
     """Auto-regressive decoding without caching.
 
-        Args:
+    Args:
-          inputs: input time-series and paddings. Time-series shape B x C, padding
+      inputs: input time-series and paddings. Time-series shape B x C, padding
-            shape shape B x (C + H) where H is the prediction length.
+        shape shape B x (C + H) where H is the prediction length.
-          horizon_len: prediction length.
+      horizon_len: prediction length.
-          output_patch_len: output length to be fetched from one step of
+      output_patch_len: output length to be fetched from one step of
-            auto-regressive decoding.
+        auto-regressive decoding.
-          max_len: maximum training context length.
+      max_len: maximum training context length.
+      return_forecast_on_context: whether to return the model forecast on the
+        context except the first input patch.
 
-        Returns:
+    Returns:
-          Tuple of two forecasting results:
+      Tuple of two forecasting results:
-          - Point (mean) output predictions as a tensor with shape B x H.
+      - Point (mean) output predictions as a tensor with shape B x H'.
-          - Full predictions (mean and quantiles) as a tensor with shape
+      - Full predictions (mean and quantiles) as a tensor with shape
-            B x H x (1 + # quantiles).
+        B x H' x (1 + # quantiles).
-        """
+      In particular, if return_forecast_on_context is True, H' is H plus
+      the forecastable context length, i.e. context_len - (first) patch_len.
+    """
     final_out = inputs[_INPUT_TS]
-    inp_time_len = final_out.shape[1]
+    context_len = final_out.shape[1]
     paddings = inputs[_INPUT_PADDING]
     if self.use_freq:
       freq = inputs[_FREQ].astype(jnp.int32)
@@ -425,13 +443,15 @@ 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 -
+    num_decode_patches = (
-                          1) // output_patch_len
+        horizon_len + output_patch_len - 1
-    for _ in range(num_decode_patches):
+    ) // output_patch_len
-      current_padding = paddings[:, 0:final_out.shape[1]]
+    for step_index in range(num_decode_patches):
+      current_padding = paddings[:, 0 : final_out.shape[1]]
       input_ts = final_out[:, -max_len:]
       input_padding = current_padding[:, -max_len:]
       model_input = NestedMap(
@@ -440,25 +460,40 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
           freq=freq,
       )
       fprop_outputs = self(model_input)[_OUTPUT_TS]
+      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 = es.jax_einshape("bnph->b(np)h", new_full_ts)
+
+        full_outputs.append(new_full_ts)
+
       # (full batch, last patch, output_patch_len, index of mean forecast = 0)
       new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
+      new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]
       # (full batch, last patch, output_patch_len, all output indices)
-      full_outputs.append(fprop_outputs[:, -1, :output_patch_len, :])
+      full_outputs.append(new_full_ts)
       final_out = jnp.concatenate([final_out, new_ts], axis=-1)
 
-    return (
+    if return_forecast_on_context:
-        final_out[:, inp_time_len:inp_time_len + horizon_len],
+      # `full_outputs` indexing starts at after the first input patch.
-        jnp.concatenate(full_outputs, axis=1)[:, 0: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, :]
+
+    return (full_outputs[:, :, 0], full_outputs)
 
 
 class PatchedDecoderFinetuneModel(base_model.BaseModel):
   """Model class for finetuning patched time-series decoder.
 
-    Attributes:
+  Attributes:
-      core_layer_tpl: config for core layer.
+    core_layer_tpl: config for core layer.
-      freq: freq to finetune on.
+    freq: freq to finetune on.
-    """
+  """
 
   core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder)
   freq: int = 0
@@ -471,12 +506,14 @@ 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) //
+    context_pad = (
-                   input_patch_len) * input_patch_len - context_len
+        (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)],
+    input_padding = jnp.pad(
-                            constant_values=1)
+        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,
@@ -485,28 +522,30 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
     )
     return self.core_layer(new_input_batch)
 
-  def _quantile_loss(self, pred: JTensor, actual: JTensor,
+  def _quantile_loss(
-                     quantile: float) -> JTensor:
+      self, pred: JTensor, actual: JTensor, quantile: float
+  ) -> JTensor:
     """Calculates quantile loss.
 
-        Args:
+    Args:
-          pred: B x T
+      pred: B x T
-          actual: B x T
+      actual: B x T
-          quantile: quantile at which loss is computed.
+      quantile: quantile at which loss is computed.
 
-        Returns:
+    Returns:
-          per coordinate loss.
+      per coordinate loss.
-        """
+    """
     dev = actual - pred
     loss_first = dev * quantile
     loss_second = -dev * (1.0 - quantile)
     return 2 * jnp.where(loss_first >= 0, loss_first, loss_second)
 
-  def compute_loss(self, prediction_output: NestedMap,
+  def compute_loss(
-                   input_batch: NestedMap) -> Tuple[NestedMap, NestedMap]:
+      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/time_features.py

@@ -11,6 +11,7 @@
 # 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.
+
 """Directory to extract time covariates.
 
 Extract time covariates from datetime.
@@ -35,6 +36,7 @@ from pandas.tseries.offsets import Easter
 from sklearn.preprocessing import StandardScaler
 from tqdm import tqdm
 
+
 # This is 183 to cover half a year (in both directions), also for leap years
 # + 17 as Eastern can be between March, 22 - April, 25
 MAX_WINDOW = 183 + 17
@@ -48,8 +50,9 @@ def _distance_to_holiday(holiday):
         index - pd.Timedelta(days=MAX_WINDOW),
         index + pd.Timedelta(days=MAX_WINDOW),
     )
-    assert (len(holiday_date) != 0  # pylint: disable=g-explicit-length-test
+    assert (
-           ), f"No closest holiday for the date index {index} found."
+        len(holiday_date) != 0  # pylint: disable=g-explicit-length-test
+    ), f"No closest holiday for the date index {index} found."
     # It sometimes returns two dates if it is exactly half a year after the
     # holiday. In this case, the smaller distance (182 days) is returned.
     return (index - holiday_date[0]).days
@@ -57,19 +60,16 @@ def _distance_to_holiday(holiday):
   return _distance_to_day
 
 
-EasterSunday = Holiday("Easter Sunday",
+EasterSunday = Holiday(
-                       month=1,
+    "Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
-                       day=1,
+)
-                       offset=[Easter(), Day(0)])
 NewYearsDay = Holiday("New Years Day", month=1, day=1)
-SuperBowl = Holiday("Superbowl",
+SuperBowl = Holiday(
-                    month=2,
+    "Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1))
-                    day=1,
+)
-                    offset=DateOffset(weekday=SU(1)))
+MothersDay = Holiday(
-MothersDay = Holiday("Mothers Day",
+    "Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2))
-                     month=5,
+)
-                     day=1,
-                     offset=DateOffset(weekday=SU(2)))
 IndependenceDay = Holiday("Independence Day", month=7, day=4)
 ChristmasEve = Holiday("Christmas", month=12, day=24)
 ChristmasDay = Holiday("Christmas", month=12, day=25)

src/timesfm/timesfm.py

@@ -11,8 +11,10 @@
 # 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."""
 
+import collections
 import logging
 import multiprocessing
 from os import path
@@ -20,11 +22,11 @@ import time
 from typing import Any, Literal, Optional, 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 huggingface_hub import snapshot_download
 from paxml import checkpoints
 from paxml import tasks_lib
 from praxis import base_hyperparams
@@ -34,12 +36,19 @@ from praxis import py_utils
 from praxis import pytypes
 from praxis.layers import normalizations
 from praxis.layers import transformers
-from . import patched_decoder
+
 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
 
 
 def process_group(key, group, value_name, forecast_context_len):
@@ -51,16 +60,23 @@ 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") /
+  smoothed_arr = (
-                  window_size)
+      np.convolve(arr_padded, np.ones(window_size), "valid") / window_size
+  )
   return [smoothed_arr, arr - smoothed_arr]
 
 
 def freq_map(freq: str):
   """Returns the frequency map for the given frequency string."""
   freq = str.upper(freq)
-  if (freq.endswith("H") or freq.endswith("T") or freq.endswith("MIN") or
+  if (
-      freq.endswith("D") or freq.endswith("B") or freq.endswith("U")):
+      freq.endswith("H")
+      or freq.endswith("T")
+      or freq.endswith("MIN")
+      or freq.endswith("D")
+      or freq.endswith("B")
+      or freq.endswith("U")
+  ):
     return 0
   elif freq.endswith(("W", "M", "MS")):
     return 1
@@ -70,6 +86,20 @@ def freq_map(freq: str):
     raise ValueError(f"Invalid frequency: {freq}")
 
 
+# Per time series normalization: forward.
+def _normalize(batch):
+  stats = [
+      (np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch
+  ]
+  new_batch = [(x - stat[0]) / stat[1] for x, stat in zip(batch, stats)]
+  return new_batch, stats
+
+
+# Per time series normalization: inverse.
+def _renormalize(batch, stats):
+  return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)]
+
+
 class TimesFm:
   """TimesFM forecast API for inference.
 
@@ -149,6 +179,7 @@ class TimesFm:
     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.mesh_shape = [1, self.num_devices, 1]
     self.mesh_name = ["replica", "data", "mdl"]
@@ -171,7 +202,9 @@ class TimesFm:
             num_layers=num_layers,
             transformer_layer_params_tpl=pax_fiddle.Config(
                 transformers.Transformer,
-                ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
+                ln_tpl=pax_fiddle.Config(
+                    normalizations.RmsNorm,
+                ),
             ),
         ),
     )
@@ -189,38 +222,34 @@ class TimesFm:
 
   def _get_sample_inputs(self):
     return {
-        "input_ts":
+        "input_ts": jnp.zeros(
-            jnp.zeros(
+            (
-                (
+                self.per_core_batch_size,
-                    self.per_core_batch_size,
+                self.context_len + self.output_patch_len,
-                    self.context_len + self.output_patch_len,
-                ),
-                dtype=jnp.float32,
             ),
-        "input_padding":
+            dtype=jnp.float32,
-            jnp.zeros(
+        ),
-                (
+        "input_padding": jnp.zeros(
-                    self.per_core_batch_size,
+            (
-                    self.context_len + self.output_patch_len,
+                self.per_core_batch_size,
-                ),
+                self.context_len + self.output_patch_len,
-                dtype=jnp.float32,
             ),
-        "freq":
+            dtype=jnp.float32,
-            jnp.zeros(
+        ),
-                (
+        "freq": jnp.zeros(
-                    self.per_core_batch_size,
+            (
-                    1,
+                self.per_core_batch_size,
-                ),
+                1,
-                dtype=jnp.int32,
             ),
+            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.
+      checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX,
-      FLAX,
       step: int | None = None,
   ) -> None:
     """Loads a checkpoint and compiles the decoder.
@@ -240,7 +269,8 @@ class TimesFm:
     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)
+        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,
@@ -254,7 +284,8 @@ class TimesFm:
         learners=None,
     )
     self._logging(
-        f"Constructed model weights in {time.time() - start_time:.2f} seconds.")
+        f"Constructed model weights in {time.time() - start_time:.2f} seconds."
+    )
 
     # Load the model weights.
     self._logging(f"Restoring checkpoint from {checkpoint_path}.")
@@ -267,7 +298,8 @@ class TimesFm:
         step=step,
     )
     self._logging(
-        f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
+        f"Restored checkpoint in {time.time() - start_time:.2f} seconds."
+    )
     self.jit_decode()
 
   def jit_decode(self):
@@ -283,6 +315,7 @@ class TimesFm:
           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,
@@ -302,36 +335,34 @@ class TimesFm:
     with base_layer.JaxContext.new_context(hparams=self._eval_context):
       _ = self._pmapped_decode(
           NestedMap({
-              "input_ts":
+              "input_ts": jnp.zeros(
-                  jnp.zeros(
+                  (
-                      (
+                      self.num_devices,
-                          self.num_devices,
+                      self.per_core_batch_size,
-                          self.per_core_batch_size,
+                      self.context_len,
-                          self.context_len,
-                      ),
-                      dtype=jnp.float32,
                   ),
-              "input_padding":
+                  dtype=jnp.float32,
-                  jnp.zeros(
+              ),
-                      (
+              "input_padding": jnp.zeros(
-                          self.num_devices,
+                  (
-                          self.per_core_batch_size,
+                      self.num_devices,
-                          self.context_len + self.horizon_len,
+                      self.per_core_batch_size,
-                      ),
+                      self.context_len + self.horizon_len,
-                      dtype=jnp.float32,
                   ),
-              "date_features":
+                  dtype=jnp.float32,
-                  None,
+              ),
-              "freq":
+              "date_features": None,
-                  jnp.zeros(
+              "freq": jnp.zeros(
-                      (self.num_devices, self.per_core_batch_size, 1),
+                  (self.num_devices, self.per_core_batch_size, 1),
-                      dtype=jnp.int32,
+                  dtype=jnp.int32,
-                  ),
+              ),
-          }))
+          })
+      )
     self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
 
-  def _preprocess(self, inputs: Sequence[np.array],
+  def _preprocess(
-                  freq: Sequence[int]) -> tuple[np.array, np.array, int]:
+      self, inputs: Sequence[np.array], freq: Sequence[int]
+  ) -> tuple[np.array, np.array, int]:
     """Formats and pads raw inputs to feed into the model.
 
     This function both pads each time series to match the context length, and
@@ -352,21 +383,24 @@ class TimesFm:
 
     input_ts, input_padding, inp_freq = [], [], []
 
-    pmap_pad = ((len(inputs) - 1) // self.global_batch_size +
+    pmap_pad = (
-                1) * self.global_batch_size - len(inputs)
+        (len(inputs) - 1) // self.global_batch_size + 1
+    ) * self.global_batch_size - len(inputs)
 
     for i, ts in enumerate(inputs):
       input_len = ts.shape[0]
       padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float)
       if input_len < self.context_len:
         num_front_pad = self.context_len - input_len
-        ts = np.concatenate([np.zeros(shape=(num_front_pad,), dtype=float), ts],
+        ts = np.concatenate(
-                            axis=0)
+            [np.zeros(shape=(num_front_pad,), dtype=float), ts], axis=0
+        )
         padding = np.concatenate(
-            [np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0)
+            [np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0
+        )
       elif input_len > self.context_len:
-        ts = ts[-self.context_len:]
+        ts = ts[-self.context_len :]
-        padding = padding[-(self.context_len + self.horizon_len):]
+        padding = padding[-(self.context_len + self.horizon_len) :]
 
       input_ts.append(ts)
       input_padding.append(padding)
@@ -391,6 +425,7 @@ class TimesFm:
       freq: Sequence[int] | None = None,
       window_size: int | None = None,
       forecast_context_len: int | None = None,
+      return_forecast_on_context: bool = False,
   ) -> tuple[JTensor, JTensor]:
     """Forecasts on a list of time series.
 
@@ -403,6 +438,8 @@ class TimesFm:
       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.
 
     Returns:
     A tuple for JTensors:
@@ -416,7 +453,8 @@ class TimesFm:
     if not self._train_state or not self._model:
       raise ValueError(
           "Checkpoint not loaded. Call `load_from_checkpoint` before"
-          " `forecast`.")
+          " `forecast`."
+      )
     if forecast_context_len is None:
       forecast_context_len = self.context_len
     inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
@@ -438,45 +476,50 @@ class TimesFm:
       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) *
+        input_ts_in = jnp.array(
-                                         self.global_batch_size])
+            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) *
+            input_padding[
-                          self.global_batch_size],)
+                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) *
+            inp_freq[
-                     self.global_batch_size, :],
+                i * self.global_batch_size : (i + 1) * self.global_batch_size, :
+            ],
             dtype=jnp.int32,
         )
         pmapped_inputs = NestedMap({
-            "input_ts":
+            "input_ts": es.jax_einshape(
-                es.jax_einshape(
+                "(db)...->db...",
-                    "(db)...->db...",
+                input_ts_in,
-                    input_ts_in,
+                d=self.num_devices,
-                    d=self.num_devices,
+            ),
-                ),
+            "input_padding": es.jax_einshape(
-            "input_padding":
+                "(db)...->db...",
-                es.jax_einshape(
+                input_padding_in,
-                    "(db)...->db...",
+                d=self.num_devices,
-                    input_padding_in,
+            ),
-                    d=self.num_devices,
+            "date_features": None,
-                ),
+            "freq": es.jax_einshape(
-            "date_features":
+                "(db)...->db...",
-                None,
+                inp_freq_in,
-            "freq":
+                d=self.num_devices,
-                es.jax_einshape(
+            ),
-                    "(db)...->db...",
-                    inp_freq_in,
-                    d=self.num_devices,
-                ),
         })
         mean_output, full_output = self._pmapped_decode(pmapped_inputs)
-        mean_output = es.jax_einshape("db...->(db)...",
+        if not return_forecast_on_context:
-                                      mean_output,
+          mean_output = mean_output[:, :, self._horizon_start :, ...]
-                                      d=self.num_devices)
+          full_output = full_output[:, :, self._horizon_start :, ...]
-        full_output = es.jax_einshape("db...->(db)...",
+        mean_output = es.jax_einshape(
-                                      full_output,
+            "db...->(db)...", mean_output, d=self.num_devices
-                                      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)
@@ -497,6 +540,240 @@ class TimesFm:
       full_outputs = np.maximum(full_outputs, 0.0)
     return mean_outputs, full_outputs
 
+  def forecast_with_covariates(
+      self,
+      inputs: list[Sequence[float]],
+      dynamic_numerical_covariates: (
+          dict[str, Sequence[Sequence[float]]] | None
+      ) = None,
+      dynamic_categorical_covariates: (
+          dict[str, Sequence[Sequence[Category]]] | None
+      ) = None,
+      static_numerical_covariates: dict[str, Sequence[float]] | None = None,
+      static_categorical_covariates: (
+          dict[str, Sequence[Category]] | None
+      ) = None,
+      freq: Sequence[int] | None = None,
+      window_size: int | None = None,
+      forecast_context_len: int | None = None,
+      xreg_mode: XRegMode = "xreg + timesfm",
+      normalize_xreg_target_per_input: bool = True,
+      ridge: float = 0.0,
+      max_rows_per_col: int = 0,
+      force_on_cpu: bool = False,
+  ):
+    """Forecasts on a list of time series with covariates.
+
+    To optimize inference speed, avoid string valued categorical covariates.
+
+    Args:
+      inputs: A list of time series forecast contexts. Each context time series
+        should be in a format convertible to JTensor by `jnp.array`.
+      dynamic_numerical_covariates: A dict of dynamic numerical covariates.
+      dynamic_categorical_covariates: A dict of dynamic categorical covariates.
+      static_numerical_covariates: A dict of static numerical covariates.
+      static_categorical_covariates: A dict of static categorical covariates.
+      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.
+      xreg_mode: one of "xreg + timesfm" or "timesfm + xreg". "xreg + timesfm"
+        fits a model on the residuals of the TimesFM forecast. "timesfm + xreg"
+        fits a model on the targets then forecasts on the residuals via TimesFM.
+      normalize_xreg_target_per_input: whether to normalize the xreg target per
+        input in the given batch.
+      ridge: ridge penalty for the linear model.
+      max_rows_per_col: max number of rows per column for the linear model.
+      force_on_cpu: whether to force running on cpu for the linear model.
+
+    Returns:
+      A tuple of two lists. The first is the outputs of the model. The second is
+      the outputs of the xreg.
+    """
+
+    # Verify and bookkeep covariates.
+    if not (
+        dynamic_numerical_covariates
+        or dynamic_categorical_covariates
+        or static_numerical_covariates
+        or static_categorical_covariates
+    ):
+      raise ValueError(
+          "At least one of dynamic_numerical_covariates,"
+          " dynamic_categorical_covariates, static_numerical_covariates,"
+          " static_categorical_covariates must be set."
+      )
+
+    # Track the lengths of (1) each input, (2) the part that can be used in the
+    # linear model, and (3) the horizon.
+    input_lens, train_lens, test_lens = [], [], []
+
+    for i, input_ts in enumerate(inputs):
+      input_len = len(input_ts)
+      input_lens.append(input_len)
+
+      if xreg_mode == "timesfm + xreg":
+        # For fitting residuals, no TimesFM forecast on the first patch.
+        train_lens.append(max(0, input_len - self.input_patch_len))
+      elif xreg_mode == "xreg + timesfm":
+        train_lens.append(input_len)
+      else:
+        raise ValueError(f"Unsupported mode: {xreg_mode}")
+
+      if dynamic_numerical_covariates:
+        test_lens.append(
+            len(list(dynamic_numerical_covariates.values())[0][i]) - input_len
+        )
+      elif dynamic_categorical_covariates:
+        test_lens.append(
+            len(list(dynamic_categorical_covariates.values())[0][i]) - input_len
+        )
+      else:
+        test_lens.append(self.horizon_len)
+
+      if test_lens[-1] > self.horizon_len:
+        raise ValueError(
+            "Forecast requested longer horizon than the model definition "
+            f"supports: {test_lens[-1]} vs {self.horizon_len}."
+        )
+
+    # Prepare the covariates into train and test.
+    train_dynamic_numerical_covariates = collections.defaultdict(list)
+    test_dynamic_numerical_covariates = collections.defaultdict(list)
+    train_dynamic_categorical_covariates = collections.defaultdict(list)
+    test_dynamic_categorical_covariates = collections.defaultdict(list)
+    for covariates, train_covariates, test_covariates in (
+        (
+            dynamic_numerical_covariates,
+            train_dynamic_numerical_covariates,
+            test_dynamic_numerical_covariates,
+        ),
+        (
+            dynamic_categorical_covariates,
+            train_dynamic_categorical_covariates,
+            test_dynamic_categorical_covariates,
+        ),
+    ):
+      if not covariates:
+        continue
+      for covariate_name, covariate_values in covariates.items():
+        for input_len, train_len, covariate_value in zip(
+            input_lens, train_lens, covariate_values
+        ):
+          train_covariates[covariate_name].append(
+              covariate_value[(input_len - train_len) : input_len]
+          )
+          test_covariates[covariate_name].append(covariate_value[input_len:])
+
+    # Fit models.
+    if xreg_mode == "timesfm + xreg":
+      # Forecast via TimesFM then fit a model on the residuals.
+      mean_outputs, _ = self.forecast(
+          inputs,
+          freq,
+          window_size,
+          forecast_context_len,
+          return_forecast_on_context=True,
+      )
+      targets = [
+          (
+              np.array(input_ts)[-train_len:]
+              - mean_output[
+                  (self._horizon_start - train_len) : self._horizon_start
+              ]
+          )
+          for input_ts, mean_output, train_len in zip(
+              inputs, mean_outputs, train_lens
+          )
+      ]
+      per_instance_stats = None
+      if normalize_xreg_target_per_input:
+        targets, per_instance_stats = _normalize(targets)
+      xregs = xreg_lib.BatchedInContextXRegLinear(
+          targets=targets,
+          train_lens=train_lens,
+          test_lens=test_lens,
+          train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
+          test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
+          train_dynamic_categorical_covariates=train_dynamic_categorical_covariates,
+          test_dynamic_categorical_covariates=test_dynamic_categorical_covariates,
+          static_numerical_covariates=static_numerical_covariates,
+          static_categorical_covariates=static_categorical_covariates,
+      ).fit(
+          ridge=ridge,
+          one_hot_encoder_drop=None if ridge > 0 else "first",
+          max_rows_per_col=max_rows_per_col,
+          force_on_cpu=force_on_cpu,
+          debug_info=False,
+          assert_covariates=True,
+          assert_covariate_shapes=True,
+      )
+      if normalize_xreg_target_per_input:
+        xregs = _renormalize(xregs, per_instance_stats)
+      outputs = [
+          (
+              mean_output[
+                  self._horizon_start : (self._horizon_start + test_len)
+              ]
+              + xreg
+          )
+          for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
+      ]
+
+    else:
+      # Fit a model on the targets then forecast on the residuals via TimesFM.
+      targets = [
+          np.array(input_ts)[-train_len:]
+          for input_ts, train_len in zip(inputs, train_lens)
+      ]
+      per_instance_stats = None
+      if normalize_xreg_target_per_input:
+        targets, per_instance_stats = _normalize(targets)
+      xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear(
+          targets=targets,
+          train_lens=train_lens,
+          test_lens=test_lens,
+          train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
+          test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
+          train_dynamic_categorical_covariates=train_dynamic_categorical_covariates,
+          test_dynamic_categorical_covariates=test_dynamic_categorical_covariates,
+          static_numerical_covariates=static_numerical_covariates,
+          static_categorical_covariates=static_categorical_covariates,
+      ).fit(
+          ridge=ridge,
+          one_hot_encoder_drop=None if ridge > 0 else "first",
+          max_rows_per_col=max_rows_per_col,
+          force_on_cpu=force_on_cpu,
+          debug_info=True,
+          assert_covariates=True,
+          assert_covariate_shapes=True,
+      )
+      mean_outputs, _ = self.forecast(
+          [
+              target - xreg_on_context
+              for target, xreg_on_context in zip(targets, xregs_on_context)
+          ],
+          freq,
+          window_size,
+          forecast_context_len,
+          return_forecast_on_context=True,
+      )
+      outputs = [
+          (
+              mean_output[
+                  self._horizon_start : (self._horizon_start + test_len)
+              ]
+              + xreg
+          )
+          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)
+
+    return outputs, xregs
+
   def forecast_on_df(
       self,
       inputs: pd.DataFrame,
@@ -527,10 +804,14 @@ class TimesFm:
     Returns:
       Future forecasts dataframe.
     """
-    if not ("unique_id" in inputs.columns and "ds" in inputs.columns and
+    if not (
-            value_name in inputs.columns):
+        "unique_id" in inputs.columns
+        and "ds" in inputs.columns
+        and value_name in inputs.columns
+    ):
       raise ValueError(
-          f"DataFrame must have unique_id, ds and {value_name} columns.")
+          f"DataFrame must have unique_id, ds and {value_name} columns."
+      )
     if not forecast_context_len:
       forecast_context_len = self.context_len
     logging.info("Preprocessing dataframe.")
@@ -555,15 +836,17 @@ class TimesFm:
       with multiprocessing.Pool(processes=num_jobs) as pool:
         results = pool.starmap(
             process_group,
-            [(key, group, value_name, forecast_context_len)
+            [
-             for key, group in df_sorted.groupby("unique_id")],
+                (key, group, value_name, forecast_context_len)
+                for key, group in df_sorted.groupby("unique_id")
+            ],
         )
       new_inputs, uids = zip(*results)
     print("Finished preprocessing dataframe.")
     freq_inps = [freq_map(freq)] * len(new_inputs)
-    _, full_forecast = self.forecast(new_inputs,
+    _, full_forecast = self.forecast(
-                                     freq=freq_inps,
+        new_inputs, freq=freq_inps, window_size=window_size
-                                     window_size=window_size)
+    )
     print("Finished forecasting.")
     fcst_df = make_future_dataframe(
         uids=uids,
@@ -571,13 +854,16 @@ class TimesFm:
         h=self.horizon_len,
         freq=freq,
     )
-    fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
+    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,
+        fcst_df[q_col] = full_forecast[:, 0 : self.horizon_len, 1 + i].reshape(
-                                       1 + i].reshape(-1, 1)
+            -1, 1
+        )
         if q == 0.5:
           fcst_df[model_name] = fcst_df[q_col]
     logging.info("Finished creating output dataframe.")

src/timesfm/xreg_lib.py

@@ -0,0 +1,532 @@
+# 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.
+
+"""Helper functions for in-context covariates and regression."""
+
+import itertools
+import math
+from typing import Any, Iterable, Literal, Mapping, Sequence
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from sklearn import preprocessing
+
+Category = int | str
+
+_TOL = 1e-6
+XRegMode = Literal["timesfm + xreg", "xreg + timesfm"]
+
+
+def _unnest(nested: Sequence[Sequence[Any]]) -> np.ndarray:
+  return np.array(list(itertools.chain.from_iterable(nested)))
+
+
+def _repeat(elements: Iterable[Any], counts: Iterable[int]) -> np.ndarray:
+  return np.array(
+      list(
+          itertools.chain.from_iterable(map(itertools.repeat, elements, counts))
+      )
+  )
+
+
+def _to_padded_jax_array(x: np.ndarray) -> jax.Array:
+  if x.ndim == 1:
+    (i,) = x.shape
+    di = 2 ** math.ceil(math.log2(i)) - i
+    return jnp.pad(x, ((0, di),), mode="constant", constant_values=0.0)
+  elif x.ndim == 2:
+    i, j = x.shape
+    di = 2 ** math.ceil(math.log2(i)) - i
+    dj = 2 ** math.ceil(math.log2(j)) - j
+    return jnp.pad(x, ((0, di), (0, dj)), mode="constant", constant_values=0.0)
+  else:
+    raise ValueError(f"Unsupported array shape: {x.shape}")
+
+
+class BatchedInContextXRegBase:
+  """Helper class for in-context regression covariate formatting.
+
+  Attributes:
+    targets: List of targets (responses) of the in-context regression.
+    train_lens: List of lengths of each target vector from the context.
+    test_lens: List of lengths of each forecast horizon.
+    train_dynamic_numerical_covariates: Dict of covariate names mapping to the
+      dynamic numerical covariates of each forecast task on the context. Their
+      lengths should match the corresponding lengths in `train_lens`.
+    train_dynamic_categorical_covariates: Dict of covariate names mapping to the
+      dynamic categorical covariates of each forecast task on the context. Their
+      lengths should match the corresponding lengths in `train_lens`.
+    test_dynamic_numerical_covariates: Dict of covariate names mapping to the
+      dynamic numerical covariates of each forecast task on the horizon. Their
+      lengths should match the corresponding lengths in `test_lens`.
+    test_dynamic_categorical_covariates: Dict of covariate names mapping to the
+      dynamic categorical covariates of each forecast task on the horizon. Their
+      lengths should match the corresponding lengths in `test_lens`.
+    static_numerical_covariates: Dict of covariate names mapping to the static
+      numerical covariates of each forecast task.
+    static_categorical_covariates: Dict of covariate names mapping to the static
+      categorical covariates of each forecast task.
+  """
+
+  def __init__(
+      self,
+      targets: Sequence[Sequence[float]],
+      train_lens: Sequence[int],
+      test_lens: Sequence[int],
+      train_dynamic_numerical_covariates: (
+          Mapping[str, Sequence[Sequence[float]]] | None
+      ) = None,
+      train_dynamic_categorical_covariates: (
+          Mapping[str, Sequence[Sequence[Category]]] | None
+      ) = None,
+      test_dynamic_numerical_covariates: (
+          Mapping[str, Sequence[Sequence[float]]] | None
+      ) = None,
+      test_dynamic_categorical_covariates: (
+          Mapping[str, Sequence[Sequence[Category]]] | None
+      ) = None,
+      static_numerical_covariates: Mapping[str, Sequence[float]] | None = None,
+      static_categorical_covariates: (
+          Mapping[str, Sequence[Category]] | None
+      ) = None,
+  ) -> None:
+    """Initializes with the exogenous covariate inputs.
+
+    Here we use model fitting language to refer to the context as 'train' and
+    the horizon as 'test'. We assume batched inputs. To properly format the
+    request:
+
+     - `train_lens` represents the contexts in the batch. Targets and all train
+     dynamic covariates should have the same lengths as the corresponding
+     elements
+     in `train_lens`. Notice each `train_len` can be different from the exact
+     length of the corresponding context depending on how much of the context is
+     used for fitting the in-context model.
+     - `test_lens` represents the horizon lengths in the batch. All tesdt
+     dynamic
+     covariates should have the same lengths as the corresponding elements in
+     `test_lens`.
+     - Static covariates should be one for each input.
+     - For train and test dynamic covariates, they should have the same
+     covariate
+     names.
+
+     Pass an empty dict {} for a covariate type if it is not present.
+
+     Example:
+       Here is a set of valid inputs whose schema can be used for reference.
+       ```
+       targets = [
+           [0.0, 0.1, 0.2],
+           [0.0, 0.1, 0.2, 0.3],
+       ]  # Two inputs in this batch.
+       train_lens = [3, 4]
+       test_lens = [2, 5]  # Forecast horizons 2 and 5 respectively.
+       train_dynamic_numerical_covariates = {
+           "cov_1_dn": [[0.0, 0.5, 1.0], [0.0, 0.5, 1.0, 1.5]],
+           "cov_2_dn": [[0.0, 1.5, 1.0], [0.0, 1.5, 1.0, 2.5]],
+       }  # Each train dynamic covariate has 3 and 4 elements respectively.
+       test_dynamic_numerical_covariates = {
+           "cov_1_dn": [[0.1, 0.6], [0.1, 0.6, 1.1, 1.6, 2.4]],
+           "cov_2_dn": [[0.1, 1.1], [0.1, 1.6, 1.1, 2.6, 10.0]],
+       }  # Each test dynamic covariate has 2 and 5 elements respectively.
+       train_dynamic_categorical_covariates = {
+           "cov_1_dc": [[0, 1, 0], [0, 1, 2, 3]],
+           "cov_2_dc": [["good", "bad", "good"], ["good", "good", "bad",
+           "bad"]],
+       }
+       test_dynamic_categorical_covariates = {
+           "cov_1_dc": [[1, 0], [1, 0, 2, 3, 1]],
+           "cov_2_dc": [["bad", "good"], ["bad", "bad", "bad", "bad", "bad"]],
+       }
+       static_numerical_covariates = {
+           "cov_1_sn": [0.0, 3.0],
+           "cov_2_sn": [2.0, 1.0],
+           "cov_3_sn": [1.0, 2.0],
+       }  # Each static covariate has 1 element for each input.
+       static_categorical_covariates = {
+           "cov_1_sc": ["apple", "orange"],
+           "cov_2_sc": [2, 3],
+       }
+       ```
+
+    Args:
+      targets: List of targets (responses) of the in-context regression.
+      train_lens: List of lengths of each target vector from the context.
+      test_lens: List of lengths of each forecast horizon.
+      train_dynamic_numerical_covariates: Dict of covariate names mapping to the
+        dynamic numerical covariates of each forecast task on the context. Their
+        lengths should match the corresponding lengths in `train_lens`.
+      train_dynamic_categorical_covariates: Dict of covariate names mapping to
+        the dynamic categorical covariates of each forecast task on the context.
+        Their lengths should match the corresponding lengths in `train_lens`.
+      test_dynamic_numerical_covariates: Dict of covariate names mapping to the
+        dynamic numerical covariates of each forecast task on the horizon. Their
+        lengths should match the corresponding lengths in `test_lens`.
+      test_dynamic_categorical_covariates: Dict of covariate names mapping to
+        the dynamic categorical covariates of each forecast task on the horizon.
+        Their lengths should match the corresponding lengths in `test_lens`.
+      static_numerical_covariates: Dict of covariate names mapping to the static
+        numerical covariates of each forecast task.
+      static_categorical_covariates: Dict of covariate names mapping to the
+        static categorical covariates of each forecast task.
+    """
+    self.targets = targets
+    self.train_lens = train_lens
+    self.test_lens = test_lens
+    self.train_dynamic_numerical_covariates = (
+        train_dynamic_numerical_covariates or {}
+    )
+    self.train_dynamic_categorical_covariates = (
+        train_dynamic_categorical_covariates or {}
+    )
+    self.test_dynamic_numerical_covariates = (
+        test_dynamic_numerical_covariates or {}
+    )
+    self.test_dynamic_categorical_covariates = (
+        test_dynamic_categorical_covariates or {}
+    )
+    self.static_numerical_covariates = static_numerical_covariates or {}
+    self.static_categorical_covariates = static_categorical_covariates or {}
+
+  def _assert_covariates(self, assert_covariate_shapes: bool = False) -> None:
+    """Verifies the validity of the covariate inputs."""
+
+    # Check presence.
+    if (
+        self.train_dynamic_numerical_covariates
+        and not self.test_dynamic_numerical_covariates
+    ) or (
+        not self.train_dynamic_numerical_covariates
+        and self.test_dynamic_numerical_covariates
+    ):
+      raise ValueError(
+          "train_dynamic_numerical_covariates and"
+          " test_dynamic_numerical_covariates must be both present or both"
+          " absent."
+      )
+
+    if (
+        self.train_dynamic_categorical_covariates
+        and not self.test_dynamic_categorical_covariates
+    ) or (
+        not self.train_dynamic_categorical_covariates
+        and self.test_dynamic_categorical_covariates
+    ):
+      raise ValueError(
+          "train_dynamic_categorical_covariates and"
+          " test_dynamic_categorical_covariates must be both present or both"
+          " absent."
+      )
+
+    # Check keys.
+    for dict_a, dict_b, dict_a_name, dict_b_name in (
+        (
+            self.train_dynamic_numerical_covariates,
+            self.test_dynamic_numerical_covariates,
+            "train_dynamic_numerical_covariates",
+            "test_dynamic_numerical_covariates",
+        ),
+        (
+            self.train_dynamic_categorical_covariates,
+            self.test_dynamic_categorical_covariates,
+            "train_dynamic_categorical_covariates",
+            "test_dynamic_categorical_covariates",
+        ),
+    ):
+      if w := set(dict_a.keys()) - set(dict_b.keys()):
+        raise ValueError(
+            f"{dict_a_name} has keys not present in {dict_b_name}: {w}"
+        )
+      if w := set(dict_b.keys()) - set(dict_a.keys()):
+        raise ValueError(
+            f"{dict_b_name} has keys not present in {dict_a_name}: {w}"
+        )
+
+    # Check shapes.
+    if assert_covariate_shapes:
+      if len(self.targets) != len(self.train_lens):
+        raise ValueError(
+            "targets and train_lens must have the same number of elements."
+        )
+
+      if len(self.train_lens) != len(self.test_lens):
+        raise ValueError(
+            "train_lens and test_lens must have the same number of elements."
+        )
+
+      for i, (target, train_len) in enumerate(
+          zip(self.targets, self.train_lens)
+      ):
+        if len(target) != train_len:
+          raise ValueError(
+              f"targets[{i}] has length {len(target)} != expected {train_len}."
+          )
+
+      for key, values in self.static_numerical_covariates.items():
+        if len(values) != len(self.train_lens):
+          raise ValueError(
+              f"static_numerical_covariates has key {key} with number of"
+              f" examples {len(values)} != expected {len(self.train_lens)}."
+          )
+
+      for key, values in self.static_categorical_covariates.items():
+        if len(values) != len(self.train_lens):
+          raise ValueError(
+              f"static_categorical_covariates has key {key} with number of"
+              f" examples {len(values)} != expected {len(self.train_lens)}."
+          )
+
+      for lens, dict_cov, dict_cov_name in (
+          (
+              self.train_lens,
+              self.train_dynamic_numerical_covariates,
+              "train_dynamic_numerical_covariates",
+          ),
+          (
+              self.train_lens,
+              self.train_dynamic_categorical_covariates,
+              "train_dynamic_categorical_covariates",
+          ),
+          (
+              self.test_lens,
+              self.test_dynamic_numerical_covariates,
+              "test_dynamic_numerical_covariates",
+          ),
+          (
+              self.test_lens,
+              self.test_dynamic_categorical_covariates,
+              "test_dynamic_categorical_covariates",
+          ),
+      ):
+        for key, cov_values in dict_cov.items():
+          if len(cov_values) != len(lens):
+            raise ValueError(
+                f"{dict_cov_name} has key {key} with number of examples"
+                f" {len(cov_values)} != expected {len(lens)}."
+            )
+          for i, cov_value in enumerate(cov_values):
+            if len(cov_value) != lens[i]:
+              raise ValueError(
+                  f"{dict_cov_name} has key {key} with its {i}-th example"
+                  f" length {len(cov_value)} != expected {lens[i]}."
+              )
+
+  def create_covariate_matrix(
+      self,
+      one_hot_encoder_drop: str | None = "first",
+      use_intercept: bool = True,
+      assert_covariates: bool = False,
+      assert_covariate_shapes: bool = False,
+  ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Creates target vector and covariate matrices for in context regression.
+
+    Here we use model fitting language to refer to the context as 'train' and
+    the horizon as 'test'.
+
+    Args:
+      one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
+      use_intercept: Whether to prepare an intercept (all 1) column in the
+        matrices.
+      assert_covariates: Whether to assert the validity of the covariate inputs.
+      assert_covariate_shapes: Whether to assert the shapes of the covariate
+        inputs when `assert_covariates` is True.
+
+    Returns:
+      A tuple of the target vector, the covariate matrix for the context, and
+      the covariate matrix for the horizon.
+    """
+    if assert_covariates:
+      self._assert_covariates(assert_covariate_shapes)
+
+    x_train, x_test = [], []
+
+    # Numerical features.
+    for name in sorted(self.train_dynamic_numerical_covariates):
+      x_train.append(
+          _unnest(self.train_dynamic_numerical_covariates[name])[:, np.newaxis]
+      )
+      x_test.append(
+          _unnest(self.test_dynamic_numerical_covariates[name])[:, np.newaxis]
+      )
+
+    for covs in self.static_numerical_covariates.values():
+      x_train.append(_repeat(covs, self.train_lens)[:, np.newaxis])
+      x_test.append(_repeat(covs, self.test_lens)[:, np.newaxis])
+
+    if x_train:
+      x_train = np.concatenate(x_train, axis=1)
+      x_test = np.concatenate(x_test, axis=1)
+
+      # Normalize for robustness.
+      x_mean = np.mean(x_train, axis=0, keepdims=True)
+      x_std = np.where(
+          (w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0
+      )
+      x_train = [(x_train - x_mean) / x_std]
+      x_test = [(x_test - x_mean) / x_std]
+
+    # Categorical features. Encode one by one.
+    one_hot_encoder = preprocessing.OneHotEncoder(
+        drop=one_hot_encoder_drop,
+        sparse=False,
+        handle_unknown="ignore",
+    )
+    for name in sorted(self.train_dynamic_categorical_covariates.keys()):
+      ohe_train = _unnest(self.train_dynamic_categorical_covariates[name])[
+          :, np.newaxis
+      ]
+      ohe_test = _unnest(self.test_dynamic_categorical_covariates[name])[
+          :, np.newaxis
+      ]
+      x_train.append(np.array(one_hot_encoder.fit_transform(ohe_train)))
+      x_test.append(np.array(one_hot_encoder.transform(ohe_test)))
+
+    for covs in self.static_categorical_covariates.values():
+      ohe = one_hot_encoder.fit_transform(np.array(covs)[:, np.newaxis])
+      x_train.append(_repeat(ohe, self.train_lens))
+      x_test.append(_repeat(ohe, self.test_lens))
+
+    x_train = np.concatenate(x_train, axis=1)
+    x_test = np.concatenate(x_test, axis=1)
+
+    if use_intercept:
+      x_train = np.pad(x_train, ((0, 0), (1, 0)), constant_values=1.0)
+      x_test = np.pad(x_test, ((0, 0), (1, 0)), constant_values=1.0)
+
+    return _unnest(self.targets), x_train, x_test
+
+  def fit(self) -> Any:
+    raise NotImplementedError("Fit is not implemented.")
+
+
+class BatchedInContextXRegLinear(BatchedInContextXRegBase):
+  """Linear in-context regression model."""
+
+  def fit(
+      self,
+      ridge: float = 0.0,
+      one_hot_encoder_drop: str | None = "first",
+      use_intercept: bool = True,
+      force_on_cpu: bool = False,
+      max_rows_per_col: int = 0,
+      max_rows_per_col_sample_seed: int = 42,
+      debug_info: bool = False,
+      assert_covariates: bool = False,
+      assert_covariate_shapes: bool = False,
+  ) -> (
+      list[np.ndarray]
+      | tuple[
+          list[np.ndarray], list[np.ndarray], jax.Array, jax.Array, jax.Array
+      ]
+  ):
+    """Fits a linear model for in-context regression.
+
+    Args:
+      ridge: A non-negative value for specifying the ridge regression penalty.
+        If 0 is provided, fallback to ordinary least squares. Note this penalty
+        is added to the normalized covariate matrix.
+      one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
+      use_intercept: Whether to prepare an intercept (all 1) column in the
+        matrices.
+      force_on_cpu: Whether to force execution on cpu for accelerator machines.
+      max_rows_per_col: How many rows to subsample per column. 0 for no
+        subsampling. This is for speeding up model fitting.
+      max_rows_per_col_sample_seed: The seed for the subsampling if needed by
+        `max_rows_per_col`.
+      debug_info: Whether to return debug info.
+      assert_covariates: Whether to assert the validity of the covariate inputs.
+      assert_covariate_shapes: Whether to assert the shapes of the covariate
+        inputs when `assert_covariates` is True.
+
+    Returns:
+      If `debug_info` is False:
+        The linear fits on the horizon.
+      If `debug_info` is True:
+        A tuple of:
+        - the linear fits on the horizon,
+        - the linear fits on the context,
+        - the flattened target vector,
+        - the covariate matrix for the context, and
+        - the covariate matrix for the horizon.
+    """
+    flat_targets, x_train_raw, x_test = self.create_covariate_matrix(
+        one_hot_encoder_drop=one_hot_encoder_drop,
+        use_intercept=use_intercept,
+        assert_covariates=assert_covariates,
+        assert_covariate_shapes=assert_covariate_shapes,
+    )
+
+    x_train = x_train_raw.copy()
+    if max_rows_per_col:
+      nrows, ncols = x_train.shape
+      if nrows > (w := ncols * max_rows_per_col):
+        subsample = jax.random.choice(
+            jax.random.PRNGKey(max_rows_per_col_sample_seed),
+            nrows,
+            (w,),
+            replace=False,
+        )
+        x_train = x_train[subsample]
+        flat_targets = flat_targets[subsample]
+
+    device = jax.devices("cpu")[0] if force_on_cpu else None
+    # Runs jitted version of the solvers which are quicker at the cost of
+    # running jitting during the first time calling. Re-jitting happens whenever
+    # new (padded) shapes are encountered.
+    # Ocassionally it helps with the speed and the accuracy if we force single
+    # thread execution on cpu for accelerator machines:
+    # 1. Avoid moving data to accelarator memory.
+    # 2. Avoid precision loss if any.
+    with jax.default_device(device):
+      x_train_raw = _to_padded_jax_array(x_train_raw)
+      x_train = _to_padded_jax_array(x_train)
+      flat_targets = _to_padded_jax_array(flat_targets)
+      x_test = _to_padded_jax_array(x_test)
+      beta_hat = (
+          jnp.linalg.pinv(
+              x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]),
+              hermitian=True,
+          )
+          @ x_train.T
+          @ flat_targets
+      )
+      y_hat = x_test @ beta_hat
+      y_hat_context = x_train_raw @ beta_hat if debug_info else None
+
+    outputs = []
+    outputs_context = []
+
+    # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits.
+    train_index, test_index = 0, 0
+    for train_index_delta, test_index_delta in zip(
+        self.train_lens, self.test_lens
+    ):
+      outputs.append(
+          np.array(y_hat[test_index : (test_index + test_index_delta)])
+      )
+      if debug_info:
+        outputs_context.append(
+            np.array(
+                y_hat_context[train_index : (train_index + train_index_delta)]
+            )
+        )
+      train_index += train_index_delta
+      test_index += test_index_delta
+
+    if debug_info:
+      return outputs, outputs_context, flat_targets, x_train, x_test
+    else:
+      return outputs