standardizing styles accross all files.

experiments/baselines/timegpt_pipeline.py

@@ -34,7 +34,8 @@ def get_seasonality(freq: str) -> int:
   return _get_seasonality(freq, seasonalities={"D": 7})
 
 
-def maybe_convert_col_to_datetime(df: pd.DataFrame, col_name: str) -> pd.DataFrame:
+def maybe_convert_col_to_datetime(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])
@@ -62,13 +63,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, "unique_id": unique_id, "y": 0}  # Zero padding
+          "ds": padding_dates,
+          "unique_id": unique_id,
+          "y": 0
+      }  # Zero padding
                                )
 
       # 2d. Combine original and padding data, and append to the list
@@ -118,8 +120,7 @@ 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,
@@ -135,8 +136,7 @@ class Forecaster:
         raise ValueError(
             "Cross validation result produced less results than expected. "
             "Please verify that the frequency parameter (freq) matches your series' "
-                    "and that there aren't any missing periods."
+            "and that there aren't any missing periods.")
-                )
       results.append(result)
     out = vertical_concat(results)
     out = drop_index_if_pandas(out)
@@ -236,9 +236,11 @@ def run_timegpt(
     chunk_size = 5000
   else:
     chunk_size = None
-    fcsts_df = model.forecast(
+  fcsts_df = model.forecast(df=padded_train_df,
-        df=padded_train_df, h=horizon, level=level, freq=freq, chunk_size=chunk_size
+                            h=horizon,
-    )
+                            level=level,
+                            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/extended_benchmarks/run_timegpt.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 timegpt."""
 
 import os
@@ -25,7 +24,6 @@ from ..baselines.timegpt_pipeline import run_timegpt
 
 from .utils import ExperimentHandler
 
-
 dataset_names = [
     "m1_monthly",
     "m1_quarterly",
@@ -63,7 +61,6 @@ _MODEL_NAME = flags.DEFINE_string(
 )
 _SAVE_DIR = flags.DEFINE_string("save_dir", "./results", "Save directory")
 
-
 QUANTILES = list(np.arange(1, 10) / 10.0)
 
 
@@ -90,9 +87,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(
+    results = exp.evaluate_from_predictions(models=[model_name],
-            models=[model_name], fcsts_df=fcsts_df, times_df=time_df
+                                            fcsts_df=fcsts_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,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
@@ -26,7 +25,6 @@ import timesfm
 
 from .utils import ExperimentHandler
 
-
 dataset_names = [
     "m1_monthly",
     "m1_quarterly",
@@ -74,16 +72,14 @@ context_dict = {
     "m4_yearly": 64,
 }
 
-_MODEL_PATH = flags.DEFINE_string(
+_MODEL_PATH = flags.DEFINE_string("model_path", "/home/timesfm_q10_20240501",
-    "model_path", "/home/timesfm_q10_20240501", "Path to model"
+                                  "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)
 
 
@@ -127,9 +123,9 @@ def main():
     )
     total_time = time.time() - init_time
     time_df = pd.DataFrame({"time": [total_time], "model": model_name})
-    results = exp.evaluate_from_predictions(
+    results = exp.evaluate_from_predictions(models=[model_name],
-        models=[model_name], fcsts_df=fcsts_df, times_df=time_df
+                                            fcsts_df=fcsts_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,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.
-
 """Forked from https://github.com/Nixtla/nixtla/blob/main/experiments/amazon-chronos/src/utils.py."""
 
 from functools import partial
@@ -46,11 +45,9 @@ def quantile_loss(
     target_col: str = "y",
 ) -> pd.DataFrame:
   delta_y = df[models].sub(df[target_col], axis=0)
-  res = (
+  res = (np.maximum(q * delta_y,
-      np.maximum(q * delta_y, (q - 1) * delta_y)
+                    (q - 1) * delta_y).groupby(df[id_col],
-      .groupby(df[id_col], observed=True)
+                                               observed=True).mean())
-      .mean()
-  )
   res.index.name = id_col
   res = res.reset_index()
   return res
@@ -66,10 +63,8 @@ class ExperimentHandler:
       models_dir: str = "./models",
   ):
     if dataset not in gluonts_datasets:
-      raise Exception(
+      raise Exception(f"dataset {dataset} not found in gluonts "
-          f"dataset {dataset} not found in gluonts "
+                      f"available datasets: {', '.join(gluonts_datasets)}")
-          f"available datasets: {', '.join(gluonts_datasets)}"
-      )
     self.dataset = dataset
     self.quantiles = quantiles
     self.level = self._transform_quantiles_to_levels(quantiles)
@@ -80,10 +75,8 @@ class ExperimentHandler:
     gluonts_dataset = get_dataset(self.dataset)
     self.horizon = gluonts_dataset.metadata.prediction_length
     if self.horizon is None:
-      raise Exception(
+      raise Exception(f"horizon not found for dataset {self.dataset} "
-          f"horizon not found for dataset {self.dataset} "
+                      "experiment cannot be run")
-          "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
@@ -122,8 +115,7 @@ class ExperimentHandler:
 
   @staticmethod
   def _transform_quantiles_to_levels(quantiles: List[float]) -> List[int]:
-    level = [
+    level = [int(100 - 200 * q) for q in quantiles if q < 0.5
-        int(100 - 200 * q) for q in quantiles if q < 0.5
             ]  # in this case mean=mediain
     level = sorted(list(set(level)))
     return level
@@ -153,9 +145,8 @@ class ExperimentHandler:
       last_n: int | None = None,
   ) -> pd.DataFrame:
     with multiprocessing.Pool(os.cpu_count()) as pool:  # Create a process pool
-      results = pool.map(
+      results = pool.map(parallel_transform, zip(gluonts_dataset,
-          parallel_transform, zip(gluonts_dataset, repeat(last_n))
+                                                 repeat(last_n)))
-      )
     df = pd.concat(results)
     df = df.reset_index(drop=True)
     return df
@@ -177,9 +168,8 @@ 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(
+  def save_results(self, fcst_df: pd.DataFrame, total_time: float,
-      self, fcst_df: pd.DataFrame, total_time: float, model_name: str
+                   model_name: str):
-  ):
     self.save_dataframe(
         fcst_df,
         f"{model_name}-{self.dataset}-fcst.csv",
@@ -215,23 +205,21 @@ class ExperimentHandler:
     times_df = []
     for model in models:
       fcst_method_df = pd.read_csv(
-          f"{self.results_dir}/{model}-{self.dataset}-fcst.csv"
+          f"{self.results_dir}/{model}-{self.dataset}-fcst.csv").set_index(
-      ).set_index(["unique_id", "ds"])
+              ["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(
+    return self.evaluate_from_predictions(models=models,
-        models=models, fcsts_df=fcsts_df, times_df=times_df
+                                          fcsts_df=fcsts_df,
-    )
+                                          times_df=times_df)
 
-  def evaluate_from_predictions(
+  def evaluate_from_predictions(self, models: List[str], fcsts_df: pd.DataFrame,
-      self, models: List[str], fcsts_df: pd.DataFrame, times_df: pd.DataFrame
+                                times_df: pd.DataFrame) -> pd.DataFrame:
-  ) -> pd.DataFrame:
     test_df = self.test_df
     train_df = self.train_df
     test_df = test_df.merge(fcsts_df, how="left")
@@ -262,9 +250,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(
+    eval_df = eval_df.melt(id_vars="metric",
-        id_vars="metric", value_name="value", var_name="model"
+                           value_name="value",
-    )
+                           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,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
@@ -28,38 +27,28 @@ import torch
 import tqdm
 from timesfm import data_loader
 
-
 FLAGS = flags.FLAGS
 
-_BATCH_SIZE = flags.DEFINE_integer(
+_BATCH_SIZE = flags.DEFINE_integer("batch_size", 64,
-    "batch_size", 64, "Batch size for the randomly sampled batch"
+                                   "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 = flags.DEFINE_string("model_path", "./timesfm_q10_20240501",
-    "model_path", "./timesfm_q10_20240501", "The name of the dataset."
+                                  "The name of the dataset.")
-)
+_DATETIME_COL = flags.DEFINE_string("datetime_col", "date",
-_DATETIME_COL = flags.DEFINE_string(
+                                    "Column having datetime.")
-    "datetime_col", "date", "Column having datetime."
+_NUM_COV_COLS = flags.DEFINE_list("num_cov_cols", None,
-)
+                                  "Column having numerical features.")
-_NUM_COV_COLS = flags.DEFINE_list(
+_CAT_COV_COLS = flags.DEFINE_list("cat_cov_cols", None,
-    "num_cov_cols", None, "Column having numerical features."
+                                  "Column having categorical 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 = flags.DEFINE_bool("normalize", True,
-    "normalize", True, "normalize data for eval or not"
+                               "normalize data for eval or not")
-)
+_CONTEXT_LEN = flags.DEFINE_integer("context_len", 512,
-_CONTEXT_LEN = flags.DEFINE_integer(
+                                    "Length of the context window")
-    "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 = flags.DEFINE_string("results_dir", "./results/long_horizon",
-    "results_dir", "./results/long_horizon", "results directory"
+                                   "results directory")
-)
-
 
 DATA_DICT = {
     "ettm2": {
@@ -176,9 +165,8 @@ def eval():
       holiday=False,
       permute=False,
   )
-  eval_itr = dtl.tf_dataset(
+  eval_itr = dtl.tf_dataset(mode="test",
-      mode="test", shift=_PRED_LEN.value
+                            shift=_PRED_LEN.value).as_numpy_iterator()
-  ).as_numpy_iterator()
   model_path = _MODEL_PATH.value
   if model_path.startswith("amazon"):
     model = chronos.ChronosPipeline.from_pretrained(
@@ -213,9 +201,8 @@ def eval():
   for batch in tqdm.tqdm(eval_itr):
     past = batch[0]
     actuals = batch[3]
-    forecasts = get_forecasts(
+    forecasts = get_forecasts(model_path, model, past, int_freq,
-        model_path, model, past, int_freq, _PRED_LEN.value
+                              _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())

src/timesfm/__init__.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.
-
 """TimesFM init file."""
 
 from .timesfm import TimesFm, freq_map

src/timesfm/data_loader.py

@@ -11,13 +11,11 @@
 # 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
@@ -79,9 +77,8 @@ 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(
+    self.data_df.set_index(pd.DatetimeIndex(self.data_df[datetime_col]),
-        pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True
+                           inplace=True)
-    )
     self.num_cov_cols = num_cov_cols
     self.cat_cov_cols = cat_cov_cols
     self.ts_cols = ts_cols
@@ -94,11 +91,9 @@ 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
+        date_index, holiday=holiday).get_covariates()
-    ).get_covariates()
     self.hist_len = hist_len
     self.pred_len = pred_len
     self.batch_size = batch_size
@@ -253,9 +248,8 @@ class TimeSeriesdata(object):
       gen_fn = self.train_gen
     else:
       gen_fn = lambda: self.test_val_gen(mode, shift)
-    output_types = tuple(
+    output_types = tuple([tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 +
-        [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 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,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,7 +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.
@@ -238,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:
@@ -274,9 +270,8 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
 
     self.create_child(
         "position_emb",
-            pax_fiddle.Config(
+        pax_fiddle.Config(layers.PositionalEmbedding,
-                layers.PositionalEmbedding, embedding_dims=self.model_dims
+                          embedding_dims=self.model_dims),
-            ),
     )
 
     if self.use_freq:
@@ -290,25 +285,24 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
       )
 
   def transform_decode_state(
-        self, transform_fn: base_layer.DecodeStateTransformFn
+      self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
-    ) -> 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
+      self, inputs: JTensor,
-    ) -> Tuple[JTensor, Tuple[JTensor, 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)
+    outputs = jnp.where(
+        jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs)
     return outputs, (mu, sigma)
 
-    def _reverse_transform(
+  def _reverse_transform(self, outputs: JTensor,
-        self, outputs: JTensor, stats: Tuple[JTensor, JTensor]
+                         stats: Tuple[JTensor, JTensor]) -> JTensor:
-    ) -> JTensor:
     """Output is of shape [B, N, P, Q]."""
     mu, sigma = stats
     return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
@@ -323,10 +317,12 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     # Reshape into patches.
     patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len)
     input_padding = jnp.where(
-            jnp.abs(input_ts - PAD_VAL) < _TOLERANCE, 1, input_padding
+        jnp.abs(input_ts - PAD_VAL) < _TOLERANCE, 1, input_padding)
-        )
+    patched_pads = es.jax_einshape("b(np)->bnp",
-        patched_pads = es.jax_einshape("b(np)->bnp", input_padding, p=self.patch_len)
+                                   input_padding,
-        patched_inputs, stats = self._forward_transform(patched_inputs, patched_pads)
+                                   p=self.patch_len)
+    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)
@@ -355,9 +351,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(
+    output_ts = es.jax_einshape("bn(hq)->bnhq",
-            "bn(hq)->bnhq", output_ts, q=num_outputs, h=self.horizon_len
+                                output_ts,
-        )
+                                q=num_outputs,
+                                h=self.horizon_len)
     return self._reverse_transform(output_ts, stats)
 
   def __call__(self, inputs: NestedMap) -> NestedMap:
@@ -388,9 +385,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(
+    return NestedMap({
-            {_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats}
+        _OUTPUT_TOKENS: model_output,
-        )
+        _OUTPUT_TS: output_ts,
+        _STATS: stats
+    })
 
   def decode(
       self,
@@ -426,11 +425,11 @@ 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 _ in range(num_decode_patches):
       current_padding = paddings[:, 0:final_out.shape[1]]
       input_ts = final_out[:, -max_len:]
@@ -472,14 +471,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_pad = ((context_len + input_patch_len - 1) //
-            (context_len + input_patch_len - 1) // input_patch_len
+                   input_patch_len) * input_patch_len - context_len
-        ) * input_patch_len - context_len
 
     input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
-        input_padding = jnp.pad(
+    input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)],
-            input_padding, [(0, 0), (context_pad, 0)], constant_values=1
+                            constant_values=1)
-        )
     freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
     new_input_batch = NestedMap(
         input_ts=input_ts,
@@ -488,9 +485,8 @@ class PatchedDecoderFinetuneModel(base_model.BaseModel):
     )
     return self.core_layer(new_input_batch)
 
-    def _quantile_loss(
+  def _quantile_loss(self, pred: JTensor, actual: JTensor,
-        self, pred: JTensor, actual: JTensor, quantile: float
+                     quantile: float) -> JTensor:
-    ) -> JTensor:
     """Calculates quantile loss.
 
         Args:
@@ -506,9 +502,8 @@ 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(
+  def compute_loss(self, prediction_output: NestedMap,
-        self, prediction_output: NestedMap, input_batch: NestedMap
+                   input_batch: NestedMap) -> Tuple[NestedMap, 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], :]

src/timesfm/time_features.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.
-
 """Directory to extract time covariates.
 
 Extract time covariates from datetime.
@@ -36,7 +35,6 @@ 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
@@ -50,8 +48,7 @@ def _distance_to_holiday(holiday):
         index - pd.Timedelta(days=MAX_WINDOW),
         index + pd.Timedelta(days=MAX_WINDOW),
     )
-    assert (
+    assert (len(holiday_date) != 0  # pylint: disable=g-explicit-length-test
-        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.
@@ -60,16 +57,19 @@ def _distance_to_holiday(holiday):
   return _distance_to_day
 
 
-EasterSunday = Holiday(
+EasterSunday = Holiday("Easter Sunday",
-    "Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
+                       month=1,
-)
+                       day=1,
+                       offset=[Easter(), Day(0)])
 NewYearsDay = Holiday("New Years Day", month=1, day=1)
-SuperBowl = Holiday(
+SuperBowl = Holiday("Superbowl",
-    "Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1))
+                    month=2,
-)
+                    day=1,
-MothersDay = Holiday(
+                    offset=DateOffset(weekday=SU(1)))
-    "Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2))
+MothersDay = Holiday("Mothers Day",
-)
+                     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,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.
-
 """TimesFM forecast API for inference."""
 
 import logging
@@ -52,23 +51,16 @@ 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 = (
+  smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") /
-      np.convolve(arr_padded, np.ones(window_size), "valid") / window_size
+                  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 (
+  if (freq.endswith("H") or freq.endswith("T") or freq.endswith("MIN") or
-      freq.endswith("H")
+      freq.endswith("D") or freq.endswith("B") or freq.endswith("U")):
-      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
@@ -179,9 +171,7 @@ class TimesFm:
             num_layers=num_layers,
             transformer_layer_params_tpl=pax_fiddle.Config(
                 transformers.Transformer,
-                ln_tpl=pax_fiddle.Config(
+                ln_tpl=pax_fiddle.Config(normalizations.RmsNorm,),
-                    normalizations.RmsNorm,
-                ),
             ),
         ),
     )
@@ -199,21 +189,24 @@ class TimesFm:
 
   def _get_sample_inputs(self):
     return {
-        "input_ts": jnp.zeros(
+        "input_ts":
+            jnp.zeros(
                 (
                     self.per_core_batch_size,
                     self.context_len + self.output_patch_len,
                 ),
                 dtype=jnp.float32,
             ),
-        "input_padding": jnp.zeros(
+        "input_padding":
+            jnp.zeros(
                 (
                     self.per_core_batch_size,
                     self.context_len + self.output_patch_len,
                 ),
                 dtype=jnp.float32,
             ),
-        "freq": jnp.zeros(
+        "freq":
+            jnp.zeros(
                 (
                     self.per_core_batch_size,
                     1,
@@ -226,7 +219,8 @@ class TimesFm:
       self,
       checkpoint_path: Optional[str] = None,
       repo_id: str = "google/timesfm-1.0-200m",
-      checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.FLAX,
+      checkpoint_type: checkpoints.CheckpointType = checkpoints.CheckpointType.
+      FLAX,
       step: int | None = None,
   ) -> None:
     """Loads a checkpoint and compiles the decoder.
@@ -246,8 +240,7 @@ 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,
@@ -261,8 +254,7 @@ 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}.")
@@ -275,12 +267,12 @@ 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):
     """Jitting decoding function."""
+
     # Initialize and jit the decode fn.
     def _decode(inputs):
       assert self._model is not None
@@ -310,7 +302,8 @@ class TimesFm:
     with base_layer.JaxContext.new_context(hparams=self._eval_context):
       _ = self._pmapped_decode(
           NestedMap({
-              "input_ts": jnp.zeros(
+              "input_ts":
+                  jnp.zeros(
                       (
                           self.num_devices,
                           self.per_core_batch_size,
@@ -318,7 +311,8 @@ class TimesFm:
                       ),
                       dtype=jnp.float32,
                   ),
-              "input_padding": jnp.zeros(
+              "input_padding":
+                  jnp.zeros(
                       (
                           self.num_devices,
                           self.per_core_batch_size,
@@ -326,18 +320,18 @@ class TimesFm:
                       ),
                       dtype=jnp.float32,
                   ),
-              "date_features": None,
+              "date_features":
-              "freq": jnp.zeros(
+                  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 _preprocess(
+  def _preprocess(self, inputs: Sequence[np.array],
-      self, inputs: Sequence[np.array], freq: Sequence[int]
+                  freq: Sequence[int]) -> tuple[np.array, np.array, 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
@@ -358,21 +352,18 @@ class TimesFm:
 
     input_ts, input_padding, inp_freq = [], [], []
 
-    pmap_pad = (
+    pmap_pad = ((len(inputs) - 1) // self.global_batch_size +
-        (len(inputs) - 1) // self.global_batch_size + 1
+                1) * self.global_batch_size - len(inputs)
-    ) * 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(
+        ts = np.concatenate([np.zeros(shape=(num_front_pad,), dtype=float), ts],
-            [np.zeros(shape=(num_front_pad,), dtype=float), ts], axis=0
+                            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:]
         padding = padding[-(self.context_len + self.horizon_len):]
@@ -425,8 +416,7 @@ 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]
@@ -448,47 +438,45 @@ 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_in = jnp.array(input_ts[i * self.global_batch_size:(i + 1) *
-            input_ts[
+                                         self.global_batch_size])
-                i * self.global_batch_size : (i + 1) * self.global_batch_size
-            ]
-        )
         input_padding_in = jnp.array(
-            input_padding[
+            input_padding[i * self.global_batch_size:(i + 1) *
-                i * self.global_batch_size : (i + 1) * self.global_batch_size
+                          self.global_batch_size],)
-            ],
-        )
         inp_freq_in = jnp.array(
-            inp_freq[
+            inp_freq[i * self.global_batch_size:(i + 1) *
-                i * self.global_batch_size : (i + 1) * self.global_batch_size, :
+                     self.global_batch_size, :],
-            ],
             dtype=jnp.int32,
         )
         pmapped_inputs = NestedMap({
-            "input_ts": es.jax_einshape(
+            "input_ts":
+                es.jax_einshape(
                     "(db)...->db...",
                     input_ts_in,
                     d=self.num_devices,
                 ),
-            "input_padding": es.jax_einshape(
+            "input_padding":
+                es.jax_einshape(
                     "(db)...->db...",
                     input_padding_in,
                     d=self.num_devices,
                 ),
-            "date_features": None,
+            "date_features":
-            "freq": es.jax_einshape(
+                None,
+            "freq":
+                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(
+        mean_output = es.jax_einshape("db...->(db)...",
-            "db...->(db)...", mean_output, d=self.num_devices
+                                      mean_output,
-        )
+                                      d=self.num_devices)
-        full_output = es.jax_einshape(
+        full_output = es.jax_einshape("db...->(db)...",
-            "db...->(db)...", full_output, d=self.num_devices
+                                      full_output,
-        )
+                                      d=self.num_devices)
         mean_output = np.array(mean_output)
         full_output = np.array(full_output)
         mean_outputs.append(mean_output)
@@ -539,14 +527,10 @@ class TimesFm:
     Returns:
       Future forecasts dataframe.
     """
-    if not (
+    if not ("unique_id" in inputs.columns and "ds" in inputs.columns and
-        "unique_id" in inputs.columns
+            value_name 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.")
@@ -571,17 +555,15 @@ class TimesFm:
       with multiprocessing.Pool(processes=num_jobs) as pool:
         results = pool.starmap(
             process_group,
-            [
+            [(key, group, value_name, forecast_context_len)
-                (key, group, value_name, forecast_context_len)
+             for key, group in df_sorted.groupby("unique_id")],
-                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(
+    _, full_forecast = self.forecast(new_inputs,
-        new_inputs, freq=freq_inps, window_size=window_size
+                                     freq=freq_inps,
-    )
+                                     window_size=window_size)
     print("Finished forecasting.")
     fcst_df = make_future_dataframe(
         uids=uids,
@@ -589,16 +571,13 @@ class TimesFm:
         h=self.horizon_len,
         freq=freq,
     )
-    fcst_df[model_name] = full_forecast[:, 0 : self.horizon_len, 0].reshape(
+    fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
-        -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(
+        fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
-            -1, 1
+                                       1 + i].reshape(-1, 1)
-        )
         if q == 0.5:
           fcst_df[model_name] = fcst_df[q_col]
     logging.info("Finished creating output dataframe.")