2.0.0 initial

v1/src/adapter/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+
+"""adapter init file."""
+
+from .dora_layers import DoraAttentionProjection, DoraCombinedQKVProjection, DoraLinear
+from .lora_layers import LoraAttentionProjection, LoraCombinedQKVProjection, LoraLinear

v1/src/adapter/dora_layers.py

@@ -0,0 +1,202 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+
+from jax import numpy as jnp
+from praxis import base_layer
+from praxis.layers import attentions, linears
+
+WeightInit = base_layer.WeightInit
+WeightHParams = base_layer.WeightHParams
+
+
+class DoraTheta(base_layer.Theta):
+    def __init__(self, module):
+        self.module = module
+
+    def _dora_initialized(self):
+        if (
+            self.module.has_variable("params", "lora_a")
+            and self.module.has_variable("params", "lora_b")
+            and self.module.has_variable("params", "dora_m")
+            and "lora_a" in self.module._weight_hparams
+            and "lora_b" in self.module._weight_hparams
+            and "dora_m" in self.module._weight_hparams
+        ):
+            return True
+        else:
+            return False
+
+    def _dorafy_var(self, w):
+        lora_a = super().__getattr__("lora_a")
+        lora_b = super().__getattr__("lora_b")
+        dora_m = super().__getattr__("dora_m")
+
+        lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
+        lora_delta = jnp.reshape(lora_delta, w.shape)
+
+        w_prime = w + lora_delta
+
+        column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
+        norm_adapted = w_prime / column_norm
+        w_prime = dora_m * norm_adapted
+        return w_prime
+
+    def __getattr__(self, k):
+        var = super().__getattr__(k)
+        if not self._dora_initialized():
+            return var
+
+        if k == "w":
+            return self._dorafy_var(var)
+
+        return var
+
+    def __getitem__(self, k):
+        var = super().__getattr__(k)
+        if not self._dora_initialized():
+            return var
+
+        if k == "w":
+            return self._dorafy_var(var)
+
+        return var
+
+
+class DoraThetaDescriptor:
+    """Dot syntax accession descriptor."""
+
+    def __get__(self, obj, objtype=None):
+        return DoraTheta(obj)
+
+
+class DoraLinear(linears.Linear):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = DoraThetaDescriptor()
+
+    def setup(self) -> None:
+        lora_init = self.lora_init if self.lora_init else self.weight_init
+
+        super().setup()
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dims, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.output_dims, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+        self.create_variable(
+            "dora_m",
+            WeightHParams(
+                shape=[1, self.output_dims],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+
+
+class DoraAttentionProjection(attentions.AttentionProjection):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = DoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+        self.create_variable(
+            "dora_m",
+            WeightHParams(
+                shape=[1, self.num_heads, self.dim_per_head],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+
+
+class DoraCombinedQKVProjection(attentions.CombinedQKVProjectionLayer):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = DoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[3, self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[3, self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+        self.create_variable(
+            "dora_m",
+            WeightHParams(
+                shape=[3, 1, self.num_heads, self.dim_per_head],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None, None],
+            ),
+        )

v1/src/adapter/lora_layers.py

@@ -0,0 +1,166 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+
+from jax import numpy as jnp
+from praxis import base_layer
+from praxis.layers import attentions, linears
+
+WeightInit = base_layer.WeightInit
+WeightHParams = base_layer.WeightHParams
+
+
+class LoraTheta(base_layer.Theta):
+    def __init__(self, module):
+        self.module = module
+
+    def _lora_initialized(self):
+        if (
+            self.module.has_variable("params", "lora_a")
+            and self.module.has_variable("params", "lora_b")
+            and "lora_a" in self.module._weight_hparams
+            and "lora_b" in self.module._weight_hparams
+        ):
+            return True
+        else:
+            return False
+
+    def _lorafy_var(self, w):
+        lora_a = super().__getattr__("lora_a")
+        lora_b = super().__getattr__("lora_b")
+        lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
+        lora_delta = jnp.reshape(lora_delta, w.shape)
+        w_prime = w + lora_delta
+        return w_prime
+
+    def __getattr__(self, k):
+        var = super().__getattr__(k)
+        if not self._lora_initialized():
+            return var
+
+        if k == "w":
+            return self._lorafy_var(var)
+
+        return var
+
+    def __getitem__(self, k):
+        var = super().__getattr__(k)
+        if not self._lora_initialized():
+            return var
+
+        if k == "w":
+            return self._lorafy_var(var)
+
+        return var
+
+
+class LoraThetaDescriptor:
+    """Dot syntax accession descriptor."""
+
+    def __get__(self, obj, objtype=None):
+        return LoraTheta(obj)
+
+
+class LoraLinear(linears.Linear):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = LoraThetaDescriptor()
+
+    def setup(self) -> None:
+        lora_init = self.lora_init if self.lora_init else self.weight_init
+
+        super().setup()
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dims, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.output_dims, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None],
+            ),
+        )
+
+
+class LoraAttentionProjection(attentions.AttentionProjection):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = LoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[
+                    None,
+                    None,
+                ],
+            ),
+        )
+
+
+class LoraCombinedQKVProjection(attentions.CombinedQKVProjectionLayer):
+    rank: int = 0
+    lora_init: WeightInit | None = None
+    theta = LoraThetaDescriptor()
+
+    def setup(self) -> None:
+        super().setup()
+        w_weight_params = self._weight_hparams["w"]
+        lora_init = self.lora_init if self.lora_init else w_weight_params.init
+
+        self.create_variable(
+            "lora_a",
+            WeightHParams(
+                shape=[3, self.input_dim, self.rank],
+                init=lora_init,
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )
+        self.create_variable(
+            "lora_b",
+            WeightHParams(
+                shape=[3, self.dim_per_head * self.num_heads, self.rank],
+                init=WeightInit.Constant(scale=0.0),
+                mesh_shape=self.mesh_shape,
+                tensor_split_dims_mapping=[None, None, None],
+            ),
+        )

v1/src/adapter/utils.py

@@ -0,0 +1,487 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+
+"""
+This file provides functionality for loading and merging adapter weights
+in timesfm model, specifically for LoRA and DoRA.
+LoRA: https://arxiv.org/abs/2106.09685
+DoRA: https://arxiv.org/abs/2402.09353v4 
+"""
+
+import time
+
+import jax
+import jax.numpy as jnp
+from paxml import checkpoints, tasks_lib
+from paxml.train_states import TrainState
+from praxis import pax_fiddle
+
+from adapter.dora_layers import (
+    DoraAttentionProjection,
+    DoraCombinedQKVProjection,
+    DoraLinear,
+)
+from adapter.lora_layers import (
+    LoraAttentionProjection,
+    LoraCombinedQKVProjection,
+    LoraLinear,
+)
+from timesfm import TimesFm
+
+
+def get_adapter_params(
+    params: dict, lora_target_modules: str, num_layers: int, use_dora: bool = False
+) -> dict:
+    """
+    Extracts adapter parameters from the given model parameters for saving the checkpoint.
+
+    Args:
+        params (dict): The full model parameters.
+        lora_target_modules (str): Target modules for LoRA/DoRA adaptation.
+        num_layers (int): Number of transformer layers.
+        use_dora (bool, optional): Whether DoRA was used or not. Defaults to False.
+
+    Returns:
+        dict: A dictionary containing the extracted adapter parameters.
+    """
+    adapter_params = {}
+    for i in range(num_layers):
+        layer_key = f"x_layers_{i}"
+        adapter_params[layer_key] = {}
+
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                linear = params["params"]["core_layer"]["stacked_transformer_layer"][
+                    layer_key
+                ]["ff_layer"][ff_layer_key]["linear"]
+
+                lora_a = linear["lora_a"]
+                lora_b = linear["lora_b"]
+
+                adapter_params[layer_key][ff_layer_key] = {
+                    "lora_a": lora_a,
+                    "lora_b": lora_b,
+                }
+
+                if use_dora:
+                    adapter_params[layer_key][ff_layer_key]["dora_m"] = linear["dora_m"]
+
+        if lora_target_modules in ["all", "attention"]:
+            attention = params["params"]["core_layer"]["stacked_transformer_layer"][
+                layer_key
+            ]["self_attention"]
+
+            for component in ["key", "query", "value", "post"]:
+                lora_a = attention[component]["lora_a"]
+                lora_b = attention[component]["lora_b"]
+
+                adapter_params[layer_key][component] = {
+                    "lora_a": lora_a,
+                    "lora_b": lora_b,
+                }
+
+                if use_dora:
+                    adapter_params[layer_key][component]["dora_m"] = attention[
+                        component
+                    ]["dora_m"]
+    return adapter_params
+
+
+def load_adapter_checkpoint(
+    model: TimesFm,
+    adapter_checkpoint_path: str,
+    lora_rank: int,
+    lora_target_modules: str,
+    use_dora: bool,
+) -> None:
+    """
+    Loads an adapter checkpoint and merges it with the original model weights.
+
+    Args:
+        model (TimesFm): The model to update.
+        adapter_checkpoint_path (str): Path to the adapter checkpoint.
+        lora_rank (int): Rank of the LoRA adaptation.
+        lora_target_modules (str): Target modules for adaptation.
+        use_dora (bool): Whether DoRA was used or not.
+
+    Returns:
+        None
+    """
+
+    """
+    currently loading and initializing the model with adapter layers first and then merging the
+    adapter weights to original weights and replacing the adapter layers back to original layer.
+    # NOTE: refactor this. there should be a better way to load the LoRA checkpoint.
+    """
+    model._logging(f"Restoring adapter checkpoint from {adapter_checkpoint_path}.")
+    start_time = time.time()
+    original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl = (
+        load_adapter_layer(
+            mdl_vars=model._train_state.mdl_vars,
+            model=model._model,
+            lora_rank=lora_rank,
+            lora_target_modules=lora_target_modules,
+            use_dora=use_dora,
+        )
+    )
+
+    var_weight_hparams = model._model.abstract_init_with_metadata(
+        model._get_sample_inputs(), do_eval=True
+    )
+
+    adapter_weight_hparams = _get_adapter_weight_params(
+        var_weight_hparams=var_weight_hparams,
+        lora_target_modules=lora_target_modules,
+        num_layers=model._model.stacked_transformer_params_tpl.num_layers,
+        use_dora=use_dora,
+    )
+
+    adapter_state_partition_specs = tasks_lib.create_state_partition_specs(
+        adapter_weight_hparams,
+        mesh_shape=model.mesh_shape,
+        mesh_axis_names=model.mesh_name,
+        discard_opt_states=True,
+        learners=None,
+    )
+    adapter_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
+        adapter_weight_hparams,
+        discard_opt_states=True,
+        learners=None,
+    )
+    adapter_train_state = checkpoints.restore_checkpoint(
+        state_global_shapes=adapter_state_local_shapes,
+        checkpoint_dir=adapter_checkpoint_path,
+        checkpoint_type=checkpoints.CheckpointType.FLAX,
+        state_specs=adapter_state_partition_specs,
+        step=None,
+    )
+
+    # add adapter weights to the original weights
+    _merge_adapter_weights(
+        model=model,
+        adapter_train_state=adapter_train_state,
+        lora_target_modules=lora_target_modules,
+        num_layers=model._model.stacked_transformer_params_tpl.num_layers,
+        use_dora=use_dora,
+    )
+
+    # replace back with the original model layer
+    if lora_target_modules in ["all", "mlp"]:
+        model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
+            original_linear_tpl
+        )
+
+    if lora_target_modules in ["all", "attention"]:
+        model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
+            original_attn_tpl
+        )
+        model._model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
+            original_combined_qkv_tpl
+        )
+    model._logging(
+        f"Restored adapter checkpoint in {time.time() - start_time:.2f} seconds."
+    )
+
+    # jit compile the model
+    model.jit_decode()
+
+
+def _merge_adapter_weights(
+    model: TimesFm,
+    adapter_train_state: TrainState,
+    lora_target_modules: str,
+    num_layers: int,
+    use_dora: bool,
+) -> None:
+    """
+    Merges adapter weights with the original model weights.
+
+    Args:
+        model (TimesFm): The model to update.
+        adapter_train_state (TrainState): The adapter's train state.
+        lora_target_modules (str): Target modules for adaptation.
+        num_layers (int): Number of transformer layers.
+        use_dora (bool): Whether DoRA was used or not.
+    """
+    for i in range(num_layers):
+        layer_key = f"x_layers_{i}"
+
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                linear = model._train_state.mdl_vars["params"][
+                    "stacked_transformer_layer"
+                ][layer_key]["ff_layer"][ff_layer_key]["linear"]
+
+                params = adapter_train_state.mdl_vars[layer_key][ff_layer_key]
+                lora_a = params["lora_a"]
+                lora_b = params["lora_b"]
+
+                w = linear["w"]
+
+                lora_delta = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
+                lora_delta = jnp.reshape(lora_delta, w.shape)
+                w_prime = w + lora_delta
+
+                if use_dora:
+                    dora_m = params["dora_m"]
+                    column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
+                    norm_adapted = w_prime / column_norm
+                    w_prime = dora_m * norm_adapted
+                    linear["w"] = w_prime
+                    del linear["dora_m"]
+
+                else:
+                    linear["w"] = w_prime
+
+                del linear["lora_a"]
+                del linear["lora_b"]
+
+        if lora_target_modules in ["all", "attention"]:
+            attention = model._train_state.mdl_vars["params"][
+                "stacked_transformer_layer"
+            ][layer_key]["self_attention"]
+
+            for component in ["key", "query", "value", "post"]:
+                params = adapter_train_state.mdl_vars[layer_key][component]
+                lora_a = params["lora_a"]
+                lora_b = params["lora_b"]
+
+                w = attention[component]["w"]
+
+                lora_delta = jnp.einsum("...dr,...nr->...dn", lora_a, lora_b)
+                lora_delta = jnp.reshape(lora_delta, w.shape)
+                w_prime = w + lora_delta
+
+                if use_dora:
+                    dora_m = params["dora_m"]
+                    column_norm = jnp.linalg.norm(w_prime, ord=2, axis=0, keepdims=True)
+                    norm_adapted = w_prime / column_norm
+                    w_prime = dora_m * norm_adapted
+                    attention[component]["w"] = w_prime
+                    del attention[component]["dora_m"]
+
+                else:
+                    attention[component]["w"] = w_prime
+
+                del attention[component]["lora_a"]
+                del attention[component]["lora_b"]
+
+
+def _get_adapter_weight_params(
+    var_weight_hparams: dict, lora_target_modules: str, num_layers: int, use_dora: bool
+) -> dict:
+    """
+    Extracts adapter weight parameters from the given variable weight hyperparameters.
+
+    Args:
+        var_weight_hparams (dict): Variable weight hyperparameters.
+        lora_target_modules (str): Target modules for adaptation.
+        num_layers (int): Number of transformer layers.
+        use_dora (bool): Whether DoRA was used or not.
+
+    Returns:
+        dict: A dictionary containing the extracted adapter weight parameters.
+    """
+    adapter_params = {}
+    for i in range(num_layers):
+        layer = f"x_layers_{i}"
+        adapter_params[layer] = {}
+
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                adapter_weight_params = var_weight_hparams["params"][
+                    "stacked_transformer_layer"
+                ][layer]["ff_layer"][ff_layer_key]["linear"]
+                adapter_params[layer][ff_layer_key] = {
+                    "lora_a": adapter_weight_params["lora_a"],
+                    "lora_b": adapter_weight_params["lora_b"],
+                }
+
+                if use_dora:
+                    adapter_params[layer][ff_layer_key]["dora_m"] = (
+                        adapter_weight_params["dora_m"]
+                    )
+
+        if lora_target_modules in ["all", "attention"]:
+            for component in ["key", "value", "query", "post"]:
+                adapter_weight_params = var_weight_hparams["params"][
+                    "stacked_transformer_layer"
+                ][layer]["self_attention"][component]
+                adapter_params[layer][component] = {
+                    "lora_a": adapter_weight_params["lora_a"],
+                    "lora_b": adapter_weight_params["lora_b"],
+                }
+
+                if use_dora:
+                    adapter_params[layer][component]["dora_m"] = adapter_weight_params[
+                        "dora_m"
+                    ]
+
+    return adapter_params
+
+
+def load_adapter_layer(
+    mdl_vars: dict,
+    model: pax_fiddle.Config,
+    lora_rank: int,
+    lora_target_modules: str,
+    use_dora: bool = False,
+) -> tuple[pax_fiddle.Config, pax_fiddle.Config]:
+    """
+    Updates target modules with adapter layers.
+
+    Args:
+        mdl_vars (dict): Model variables.
+        model (pax_fiddle.Config): Model configuration.
+        lora_rank (int): Rank of the LoRA adaptation.
+        lora_target_modules (str): Target modules for adaptation.
+        use_dora (bool, optional): Whether DoRA was used or not.
+
+    Returns:
+        tuple[pax_fiddle.Config, pax_fiddle.Config]: Updated model configurations.
+    """
+    original_linear_tpl = original_attn_tpl = original_combined_qkv_tpl = None
+    if lora_target_modules in ["all", "mlp"]:
+        original_linear_tpl = (
+            model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl
+        )
+        adapter_linear_tpl = (
+            pax_fiddle.Config(
+                DoraLinear,
+                rank=lora_rank,
+            )
+            if use_dora
+            else pax_fiddle.Config(
+                LoraLinear,
+                rank=lora_rank,
+            )
+        )
+        adapter_linear_tpl.copy_fields_from(original_linear_tpl)
+        model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_fflayer_tpl.fflayer_tpl.linear_tpl = (
+            adapter_linear_tpl
+        )
+
+    if lora_target_modules in ["all", "attention"]:
+        original_attn_tpl = (
+            model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl
+        )
+
+        adapter_attn_tpl = (
+            pax_fiddle.Config(DoraAttentionProjection, rank=lora_rank)
+            if use_dora
+            else pax_fiddle.Config(LoraAttentionProjection, rank=lora_rank)
+        )
+        adapter_attn_tpl.copy_fields_from(original_attn_tpl)
+
+        original_combined_qkv_tpl = (
+            model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl
+        )
+
+        adapter_combined_qkv_tpl = (
+            pax_fiddle.Config(DoraCombinedQKVProjection, rank=lora_rank)
+            if use_dora
+            else pax_fiddle.Config(LoraCombinedQKVProjection, rank=lora_rank)
+        )
+        adapter_combined_qkv_tpl.copy_fields_from(original_combined_qkv_tpl)
+
+        model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.proj_tpl = (
+            adapter_attn_tpl
+        )
+        model.stacked_transformer_params_tpl.transformer_layer_params_tpl.tr_atten_tpl.combined_qkv_proj_tpl = (
+            adapter_combined_qkv_tpl
+        )
+
+    # initialize and add adapter weights
+    _initialize_adapter_params(
+        mdl_vars=mdl_vars,
+        num_layers=model.stacked_transformer_params_tpl.num_layers,
+        lora_rank=lora_rank,
+        lora_target_modules=lora_target_modules,
+        use_dora=use_dora,
+    )
+
+    return original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl
+
+
+def _initialize_adapter_params(
+    mdl_vars: dict,
+    num_layers,
+    lora_rank: int,
+    lora_target_modules: str,
+    use_dora: bool = False,
+    seed: int = 1234,
+) -> dict:
+    """
+    Initializes and adds adapter parameters to target modules.
+
+    Args:
+        mdl_vars (dict): Model variables.
+        num_layers (int): Number of transformer layers.
+        lora_rank (int): Rank of the LoRA adaptation.
+        lora_target_modules (str): Target modules for adaptation.
+        use_dora (bool, optional): Whether DoRA was used or not.
+        seed (int, optional): Random seed for initialization. Defaults to 1234.
+
+    Returns:
+        dict: Updated model variables with initialized adapter parameters.
+    """
+    for i in range(num_layers):
+        layer_key = f"x_layers_{i}"
+        if lora_target_modules in ["all", "mlp"]:
+            for ff_layer_key in ["ffn_layer1", "ffn_layer2"]:
+                linear = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
+                    "ff_layer"
+                ][ff_layer_key]["linear"]
+                original_w = linear["w"]
+                input_dim, output_dim = original_w.shape
+                std_dev = 1 / jnp.sqrt(lora_rank)
+
+                normal_initializer = jax.nn.initializers.normal(std_dev)
+                lora_a = normal_initializer(
+                    jax.random.key(seed), (input_dim, lora_rank), jnp.float32
+                )
+                lora_b = jnp.zeros((output_dim, lora_rank))
+
+                linear["lora_a"] = lora_a
+                linear["lora_b"] = lora_b
+
+                if use_dora:
+                    norm = jnp.linalg.norm(original_w, ord=2, axis=0, keepdims=True)
+                    linear["dora_m"] = norm
+
+        if lora_target_modules in ["all", "attention"]:
+            attention = mdl_vars["params"]["stacked_transformer_layer"][layer_key][
+                "self_attention"
+            ]
+
+            for component in ["key", "query", "value", "post"]:
+                original_w = attention[component]["w"]
+                w_dim = original_w.shape[0]
+                std_dev = 1 / jnp.sqrt(lora_rank)
+
+                normal_initializer = jax.nn.initializers.normal(std_dev)
+                lora_a = normal_initializer(
+                    jax.random.key(seed), (w_dim, lora_rank), jnp.float32
+                )
+                lora_b = jnp.zeros((w_dim, lora_rank))
+
+                attention[component]["lora_a"] = lora_a
+                attention[component]["lora_b"] = lora_b
+
+                if use_dora:
+                    norm = jnp.linalg.norm(
+                        original_w, ord=2, axis=0, keepdims=True
+                    ).astype(jnp.float32)
+                    attention[component]["dora_m"] = norm
+    return mdl_vars

v1/src/finetuning/finetuning_example.py

@@ -0,0 +1,403 @@
+"""
+Example usage of the TimesFM Finetuning Framework.
+
+For single GPU:
+python script.py --training_mode=single
+
+For multiple GPUs:
+python script.py --training_mode=multi --gpu_ids=0,1,2
+"""
+
+import os
+from dataclasses import asdict
+from os import path
+from typing import Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.multiprocessing as mp
+import yfinance as yf
+from absl import app, flags
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+from torch.utils.data import Dataset
+
+from finetuning.finetuning_torch import FinetuningConfig, TimesFMFinetuner
+from timesfm import TimesFm, TimesFmCheckpoint, TimesFmHparams
+from timesfm.pytorch_patched_decoder import (PatchedTimeSeriesDecoder,
+                                             TimesFMConfig)
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_enum(
+    "training_mode",
+    "single",
+    ["single", "multi"],
+    'Training mode: "single" for single-GPU or "multi" for multi-GPU training.',
+)
+
+flags.DEFINE_list(
+    "gpu_ids", ["0"],
+    "Comma-separated list of GPU IDs to use for multi-GPU training. Example: 0,1,2"
+)
+
+flags.DEFINE_string(
+    "local_model_path",
+    None,
+    "Path to a local .safetensors model file. If provided, overrides Hugging Face download."
+)
+
+class TimeSeriesDataset(Dataset):
+  """Dataset for time series data compatible with TimesFM."""
+
+  def __init__(self,
+               series: np.ndarray,
+               context_length: int,
+               horizon_length: int,
+               freq_type: int = 0):
+    """
+        Initialize dataset.
+
+        Args:
+            series: Time series data
+            context_length: Number of past timesteps to use as input
+            horizon_length: Number of future timesteps to predict
+            freq_type: Frequency type (0, 1, or 2)
+        """
+    if freq_type not in [0, 1, 2]:
+      raise ValueError("freq_type must be 0, 1, or 2")
+
+    self.series = series
+    self.context_length = context_length
+    self.horizon_length = horizon_length
+    self.freq_type = freq_type
+    self._prepare_samples()
+
+  def _prepare_samples(self) -> None:
+    """Prepare sliding window samples from the time series."""
+    self.samples = []
+    total_length = self.context_length + self.horizon_length
+
+    for start_idx in range(0, len(self.series) - total_length + 1):
+      end_idx = start_idx + self.context_length
+      x_context = self.series[start_idx:end_idx]
+      x_future = self.series[end_idx:end_idx + self.horizon_length]
+      self.samples.append((x_context, x_future))
+
+  def __len__(self) -> int:
+    return len(self.samples)
+
+  def __getitem__(
+      self, index: int
+  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    x_context, x_future = self.samples[index]
+
+    x_context = torch.tensor(x_context, dtype=torch.float32)
+    x_future = torch.tensor(x_future, dtype=torch.float32)
+
+    input_padding = torch.zeros_like(x_context)
+    freq = torch.tensor([self.freq_type], dtype=torch.long)
+
+    return x_context, input_padding, freq, x_future
+
+
+def prepare_datasets(series: np.ndarray,
+                     context_length: int,
+                     horizon_length: int,
+                     freq_type: int = 0,
+                     train_split: float = 0.8) -> Tuple[Dataset, Dataset]:
+  """
+    Prepare training and validation datasets from time series data.
+
+    Args:
+        series: Input time series data
+        context_length: Number of past timesteps to use
+        horizon_length: Number of future timesteps to predict
+        freq_type: Frequency type (0, 1, or 2)
+        train_split: Fraction of data to use for training
+
+    Returns:
+        Tuple of (train_dataset, val_dataset)
+    """
+  train_size = int(len(series) * train_split)
+  train_data = series[:train_size]
+  val_data = series[train_size:]
+
+  # Create datasets with specified frequency type
+  train_dataset = TimeSeriesDataset(train_data,
+                                    context_length=context_length,
+                                    horizon_length=horizon_length,
+                                    freq_type=freq_type)
+
+  val_dataset = TimeSeriesDataset(val_data,
+                                  context_length=context_length,
+                                  horizon_length=horizon_length,
+                                  freq_type=freq_type)
+
+  return train_dataset, val_dataset
+
+
+def get_model(load_weights: bool = False):
+  device = "cuda" if torch.cuda.is_available() else "cpu"
+  hparams = TimesFmHparams(
+      backend=device,
+      per_core_batch_size=32,
+      horizon_len=128,
+      num_layers=50,
+      use_positional_embedding=False,
+      context_len=192,
+  )
+  
+  if load_weights:
+    if FLAGS.local_model_path:
+      tfm_config = TimesFMConfig()
+      model = PatchedTimeSeriesDecoder(tfm_config)
+      loaded_checkpoint = load_file(FLAGS.local_model_path)
+    else:
+      repo_id = "google/timesfm-2.0-500m-pytorch"
+      tfm = TimesFm(hparams=hparams,
+              checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
+
+      tfm_config = tfm._model_config
+      model = PatchedTimeSeriesDecoder(tfm_config)
+      checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
+      loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+
+    model.load_state_dict(loaded_checkpoint)
+  return model, hparams, tfm_config
+
+
+def plot_predictions(
+    model: TimesFm,
+    val_dataset: Dataset,
+    save_path: Optional[str] = "predictions.png",
+) -> None:
+  """
+    Plot model predictions against ground truth for a batch of validation data.
+
+    Args:
+      model: Trained TimesFM model
+      val_dataset: Validation dataset
+      save_path: Path to save the plot
+    """
+  import matplotlib.pyplot as plt
+
+  model.eval()
+
+  x_context, x_padding, freq, x_future = val_dataset[0]
+  x_context = x_context.unsqueeze(0)  # Add batch dimension
+  x_padding = x_padding.unsqueeze(0)
+  freq = freq.unsqueeze(0)
+  x_future = x_future.unsqueeze(0)
+
+  device = next(model.parameters()).device
+  x_context = x_context.to(device)
+  x_padding = x_padding.to(device)
+  freq = freq.to(device)
+  x_future = x_future.to(device)
+
+  with torch.no_grad():
+    predictions = model(x_context, x_padding.float(), freq)
+    predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
+    last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
+
+  context_vals = x_context[0].cpu().numpy()
+  future_vals = x_future[0].cpu().numpy()
+  pred_vals = last_patch_pred[0].cpu().numpy()
+
+  context_len = len(context_vals)
+  horizon_len = len(future_vals)
+
+  plt.figure(figsize=(12, 6))
+
+  plt.plot(range(context_len),
+           context_vals,
+           label="Historical Data",
+           color="blue",
+           linewidth=2)
+
+  plt.plot(
+      range(context_len, context_len + horizon_len),
+      future_vals,
+      label="Ground Truth",
+      color="green",
+      linestyle="--",
+      linewidth=2,
+  )
+
+  plt.plot(range(context_len, context_len + horizon_len),
+           pred_vals,
+           label="Prediction",
+           color="red",
+           linewidth=2)
+
+  plt.xlabel("Time Step")
+  plt.ylabel("Value")
+  plt.title("TimesFM Predictions vs Ground Truth")
+  plt.legend()
+  plt.grid(True)
+
+  if save_path:
+    plt.savefig(save_path)
+    print(f"Plot saved to {save_path}")
+
+  plt.close()
+
+
+def get_data(context_len: int,
+             horizon_len: int,
+             freq_type: int = 0) -> Tuple[Dataset, Dataset]:
+  df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
+  time_series = df["Close"].values
+
+  train_dataset, val_dataset = prepare_datasets(
+      series=time_series,
+      context_length=context_len,
+      horizon_length=horizon_len,
+      freq_type=freq_type,
+      train_split=0.8,
+  )
+
+  print(f"Created datasets:")
+  print(f"- Training samples: {len(train_dataset)}")
+  print(f"- Validation samples: {len(val_dataset)}")
+  print(f"- Using frequency type: {freq_type}")
+  return train_dataset, val_dataset
+
+
+def single_gpu_example():
+  """Basic example of finetuning TimesFM on stock data."""
+  model, hparams, tfm_config = get_model(load_weights=True)
+  config = FinetuningConfig(batch_size=256,
+                            num_epochs=5,
+                            learning_rate=1e-4,
+                            use_wandb=True,
+                            freq_type=1,
+                            log_every_n_steps=10,
+                            val_check_interval=0.5,
+                            use_quantile_loss=True)
+
+  train_dataset, val_dataset = get_data(128,
+                                        tfm_config.horizon_len,
+                                        freq_type=config.freq_type)
+  finetuner = TimesFMFinetuner(model, config)
+
+  print("\nStarting finetuning...")
+  results = finetuner.finetune(train_dataset=train_dataset,
+                               val_dataset=val_dataset)
+
+  print("\nFinetuning completed!")
+  print(f"Training history: {len(results['history']['train_loss'])} epochs")
+
+  plot_predictions(
+      model=model,
+      val_dataset=val_dataset,
+      save_path="timesfm_predictions.png",
+  )
+
+
+def setup_process(rank, world_size, model, config, train_dataset, val_dataset,
+                  return_dict):
+  """Setup process function with optimized CUDA handling."""
+  try:
+    if torch.cuda.is_available():
+      torch.cuda.set_device(rank)
+
+    os.environ["MASTER_ADDR"] = config.master_addr
+    os.environ["MASTER_PORT"] = config.master_port
+    if not torch.distributed.is_initialized():
+      torch.distributed.init_process_group(backend="nccl",
+                                           world_size=world_size,
+                                           rank=rank)
+
+    finetuner = TimesFMFinetuner(model, config, rank=rank)
+
+    results = finetuner.finetune(train_dataset=train_dataset,
+                                 val_dataset=val_dataset)
+
+    if rank == 0:
+      return_dict["results"] = results
+      plot_predictions(
+          model=model,
+          val_dataset=val_dataset,
+          save_path="timesfm_predictions.png",
+      )
+
+  except Exception as e:
+    print(f"Error in process {rank}: {str(e)}")
+    raise e
+  finally:
+    if torch.distributed.is_initialized():
+      torch.distributed.destroy_process_group()
+
+
+def multi_gpu_example():
+  """Example of finetuning TimesFM using multiple GPUs with optimized spawn."""
+  mp.set_start_method("spawn", force=True)
+
+  gpu_ids = [0, 1]
+  world_size = len(gpu_ids)
+
+  model, hparams, tfm_config = get_model(load_weights=True)
+
+  # Create config
+  config = FinetuningConfig(
+      batch_size=256,
+      num_epochs=5,
+      learning_rate=3e-5,
+      use_wandb=True,
+      distributed=True,
+      gpu_ids=gpu_ids,
+      log_every_n_steps=50,
+      val_check_interval=0.5,
+  )
+  train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+  manager = mp.Manager()
+  return_dict = manager.dict()
+
+  # Launch processes
+  mp.spawn(
+      setup_process,
+      args=(world_size, model, config, train_dataset, val_dataset, return_dict),
+      nprocs=world_size,
+      join=True,
+  )
+
+  results = return_dict.get("results", None)
+  print("\nFinetuning completed!")
+  return results
+
+
+def main(argv):
+  """Main function that selects and runs the appropriate training mode."""
+
+  try:
+    if FLAGS.training_mode == "single":
+      print("\nStarting single-GPU training...")
+      single_gpu_example()
+    else:
+      gpu_ids = [int(id) for id in FLAGS.gpu_ids]
+      print(f"\nStarting multi-GPU training using GPUs: {gpu_ids}...")
+
+      config = FinetuningConfig(
+          batch_size=256,
+          num_epochs=5,
+          learning_rate=3e-5,
+          use_wandb=True,
+          distributed=True,
+          gpu_ids=gpu_ids,
+      )
+
+      results = multi_gpu_example(config)
+      print("\nMulti-GPU training completed!")
+
+  except Exception as e:
+    print(f"Training failed: {str(e)}")
+  finally:
+    if torch.distributed.is_initialized():
+      torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+  app.run(main)

v1/src/finetuning/finetuning_torch.py

@@ -0,0 +1,399 @@
+"""
+TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
+"""
+
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, Dataset
+from timesfm.pytorch_patched_decoder import create_quantiles
+
+import wandb
+
+
+class MetricsLogger(ABC):
+  """Abstract base class for logging metrics during training.
+
+    This class defines the interface for logging metrics during model training.
+    Concrete implementations can log to different backends (e.g., WandB, TensorBoard).
+    """
+
+  @abstractmethod
+  def log_metrics(self,
+                  metrics: Dict[str, Any],
+                  step: Optional[int] = None) -> None:
+    """Log metrics to the specified backend.
+
+        Args:
+          metrics: Dictionary containing metric names and values.
+          step: Optional step number or epoch for the metrics.
+        """
+    pass
+
+  @abstractmethod
+  def close(self) -> None:
+    """Clean up any resources used by the logger."""
+    pass
+
+
+class WandBLogger(MetricsLogger):
+  """Weights & Biases implementation of metrics logging.
+
+    Args:
+      project: Name of the W&B project.
+      config: Configuration dictionary to log.
+      rank: Process rank in distributed training.
+    """
+
+  def __init__(self, project: str, config: Dict[str, Any], rank: int = 0):
+    self.rank = rank
+    if rank == 0:
+      wandb.init(project=project, config=config)
+
+  def log_metrics(self,
+                  metrics: Dict[str, Any],
+                  step: Optional[int] = None) -> None:
+    """Log metrics to W&B if on the main process.
+
+        Args:
+          metrics: Dictionary of metrics to log.
+          step: Current training step or epoch.
+        """
+    if self.rank == 0:
+      wandb.log(metrics, step=step)
+
+  def close(self) -> None:
+    """Finish the W&B run if on the main process."""
+    if self.rank == 0:
+      wandb.finish()
+
+
+class DistributedManager:
+  """Manages distributed training setup and cleanup.
+
+    Args:
+      world_size: Total number of processes.
+      rank: Process rank.
+      master_addr: Address of the master process.
+      master_port: Port for distributed communication.
+      backend: PyTorch distributed backend to use.
+    """
+
+  def __init__(
+      self,
+      world_size: int,
+      rank: int,
+      master_addr: str = "localhost",
+      master_port: str = "12358",
+      backend: str = "nccl",
+  ):
+    self.world_size = world_size
+    self.rank = rank
+    self.master_addr = master_addr
+    self.master_port = master_port
+    self.backend = backend
+
+  def setup(self) -> None:
+    """Initialize the distributed environment."""
+    os.environ["MASTER_ADDR"] = self.master_addr
+    os.environ["MASTER_PORT"] = self.master_port
+
+    if not dist.is_initialized():
+      dist.init_process_group(backend=self.backend,
+                              world_size=self.world_size,
+                              rank=self.rank)
+
+  def cleanup(self) -> None:
+    """Clean up the distributed environment."""
+    if dist.is_initialized():
+      dist.destroy_process_group()
+
+
+@dataclass
+class FinetuningConfig:
+  """Configuration for model training.
+
+    Args:
+      batch_size: Number of samples per batch.
+      num_epochs: Number of training epochs.
+      learning_rate: Initial learning rate.
+      weight_decay: L2 regularization factor.
+      freq_type: Frequency, can be [0, 1, 2].
+      use_quantile_loss: bool = False  # Flag to enable/disable quantile loss
+      quantiles: Optional[List[float]] = None
+      device: Device to train on ('cuda' or 'cpu').
+      distributed: Whether to use distributed training.
+      gpu_ids: List of GPU IDs to use.
+      master_port: Port for distributed training.
+      master_addr: Address for distributed training.
+      use_wandb: Whether to use Weights & Biases logging.
+      wandb_project: W&B project name.
+      log_every_n_steps: Log metrics every N steps (batches), this is inspired from Pytorch Lightning
+      val_check_interval: How often within one training epoch to check val metrics. (also from Pytorch Lightning)
+        Can be: float (0.0-1.0): fraction of epoch (e.g., 0.5 = validate twice per epoch)
+                int: validate every N batches
+    """
+
+  batch_size: int = 32
+  num_epochs: int = 20
+  learning_rate: float = 1e-4
+  weight_decay: float = 0.01
+  freq_type: int = 0
+  use_quantile_loss: bool = False
+  quantiles: Optional[List[float]] = None
+  device: str = "cuda" if torch.cuda.is_available() else "cpu"
+  distributed: bool = False
+  gpu_ids: List[int] = field(default_factory=lambda: [0])
+  master_port: str = "12358"
+  master_addr: str = "localhost"
+  use_wandb: bool = False
+  wandb_project: str = "timesfm-finetuning"
+  log_every_n_steps: int = 50
+  val_check_interval: float = 0.5
+
+
+class TimesFMFinetuner:
+  """Handles model training and validation.
+
+    Args:
+      model: PyTorch model to train.
+      config: Training configuration.
+      rank: Process rank for distributed training.
+      loss_fn: Loss function (defaults to MSE).
+      logger: Optional logging.Logger instance.
+    """
+
+  def __init__(
+      self,
+      model: nn.Module,
+      config: FinetuningConfig,
+      rank: int = 0,
+      loss_fn: Optional[Callable] = None,
+      logger: Optional[logging.Logger] = None,
+  ):
+    self.model = model
+    self.config = config
+    self.rank = rank
+    self.logger = logger or logging.getLogger(__name__)
+    self.device = torch.device(
+        f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
+    self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1))**2))
+
+    if config.use_wandb:
+      self.metrics_logger = WandBLogger(config.wandb_project, config.__dict__,
+                                        rank)
+
+    if config.distributed:
+      self.dist_manager = DistributedManager(
+          world_size=len(config.gpu_ids),
+          rank=rank,
+          master_addr=config.master_addr,
+          master_port=config.master_port,
+      )
+      self.dist_manager.setup()
+      self.model = self._setup_distributed_model()
+
+  def _setup_distributed_model(self) -> nn.Module:
+    """Configure model for distributed training."""
+    self.model = self.model.to(self.device)
+    return DDP(self.model,
+               device_ids=[self.config.gpu_ids[self.rank]],
+               output_device=self.config.gpu_ids[self.rank])
+
+  def _create_dataloader(self, dataset: Dataset, is_train: bool) -> DataLoader:
+    """Create appropriate DataLoader based on training configuration.
+
+        Args:
+          dataset: Dataset to create loader for.
+          is_train: Whether this is for training (affects shuffling).
+
+        Returns:
+          DataLoader instance.
+        """
+    if self.config.distributed:
+      sampler = torch.utils.data.distributed.DistributedSampler(
+          dataset,
+          num_replicas=len(self.config.gpu_ids),
+          rank=dist.get_rank(),
+          shuffle=is_train)
+    else:
+      sampler = None
+
+    return DataLoader(
+        dataset,
+        batch_size=self.config.batch_size,
+        shuffle=(is_train and not self.config.distributed),
+        sampler=sampler,
+    )
+
+  def _quantile_loss(self, pred: torch.Tensor, actual: torch.Tensor,
+                     quantile: float) -> torch.Tensor:
+    """Calculates quantile loss.
+        Args:
+            pred: Predicted values
+            actual: Actual values
+            quantile: Quantile at which loss is computed
+        Returns:
+            Quantile loss
+        """
+    dev = actual - pred
+    loss_first = dev * quantile
+    loss_second = -dev * (1.0 - quantile)
+    return 2 * torch.where(loss_first >= 0, loss_first, loss_second)
+
+  def _process_batch(self, batch: List[torch.Tensor]) -> tuple:
+    """Process a single batch of data.
+
+        Args:
+          batch: List of input tensors.
+
+        Returns:
+          Tuple of (loss, predictions).
+        """
+    x_context, x_padding, freq, x_future = [
+        t.to(self.device, non_blocking=True) for t in batch
+    ]
+
+    predictions = self.model(x_context, x_padding.float(), freq)
+    predictions_mean = predictions[..., 0]
+    last_patch_pred = predictions_mean[:, -1, :]
+
+    loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+    if self.config.use_quantile_loss:
+      quantiles = self.config.quantiles or create_quantiles()
+      for i, quantile in enumerate(quantiles):
+        last_patch_quantile = predictions[:, -1, :, i + 1]
+        loss += torch.mean(
+            self._quantile_loss(last_patch_quantile, x_future.squeeze(-1),
+                                quantile))
+
+    return loss, predictions
+
+  def _train_epoch(self, train_loader: DataLoader,
+                   optimizer: torch.optim.Optimizer) -> float:
+    """Train for one epoch in a distributed setting.
+
+        Args:
+            train_loader: DataLoader for training data.
+            optimizer: Optimizer instance.
+
+        Returns:
+            Average training loss for the epoch.
+        """
+    self.model.train()
+    total_loss = 0.0
+    num_batches = len(train_loader)
+
+    for batch in train_loader:
+      loss, _ = self._process_batch(batch)
+
+      optimizer.zero_grad()
+      loss.backward()
+      optimizer.step()
+
+      total_loss += loss.item()
+
+    avg_loss = total_loss / num_batches
+
+    if self.config.distributed:
+      avg_loss_tensor = torch.tensor(avg_loss, device=self.device)
+      dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
+      avg_loss = (avg_loss_tensor / dist.get_world_size()).item()
+
+    return avg_loss
+
+  def _validate(self, val_loader: DataLoader) -> float:
+    """Perform validation.
+
+        Args:
+            val_loader: DataLoader for validation data.
+
+        Returns:
+            Average validation loss.
+        """
+    self.model.eval()
+    total_loss = 0.0
+    num_batches = len(val_loader)
+
+    with torch.no_grad():
+      for batch in val_loader:
+        loss, _ = self._process_batch(batch)
+        total_loss += loss.item()
+
+    avg_loss = total_loss / num_batches
+
+    if self.config.distributed:
+      avg_loss_tensor = torch.tensor(avg_loss, device=self.device)
+      dist.all_reduce(avg_loss_tensor, op=dist.ReduceOp.SUM)
+      avg_loss = (avg_loss_tensor / dist.get_world_size()).item()
+
+    return avg_loss
+
+  def finetune(self, train_dataset: Dataset,
+               val_dataset: Dataset) -> Dict[str, Any]:
+    """Train the model.
+
+        Args:
+          train_dataset: Training dataset.
+          val_dataset: Validation dataset.
+
+        Returns:
+          Dictionary containing training history.
+        """
+    self.model = self.model.to(self.device)
+    train_loader = self._create_dataloader(train_dataset, is_train=True)
+    val_loader = self._create_dataloader(val_dataset, is_train=False)
+
+    optimizer = torch.optim.Adam(self.model.parameters(),
+                                 lr=self.config.learning_rate,
+                                 weight_decay=self.config.weight_decay)
+
+    history = {"train_loss": [], "val_loss": [], "learning_rate": []}
+
+    self.logger.info(
+        f"Starting training for {self.config.num_epochs} epochs...")
+    self.logger.info(f"Training samples: {len(train_dataset)}")
+    self.logger.info(f"Validation samples: {len(val_dataset)}")
+
+    try:
+      for epoch in range(self.config.num_epochs):
+        train_loss = self._train_epoch(train_loader, optimizer)
+        val_loss = self._validate(val_loader)
+        current_lr = optimizer.param_groups[0]["lr"]
+
+        metrics = {
+            "train_loss": train_loss,
+            "val_loss": val_loss,
+            "learning_rate": current_lr,
+            "epoch": epoch + 1,
+        }
+
+        if self.config.use_wandb:
+          self.metrics_logger.log_metrics(metrics)
+
+        history["train_loss"].append(train_loss)
+        history["val_loss"].append(val_loss)
+        history["learning_rate"].append(current_lr)
+
+        if self.rank == 0:
+          self.logger.info(
+              f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}"
+          )
+
+    except KeyboardInterrupt:
+      self.logger.info("Training interrupted by user")
+
+    if self.config.distributed:
+      self.dist_manager.cleanup()
+
+    if self.config.use_wandb:
+      self.metrics_logger.close()
+
+    return {"history": history}

v1/src/timesfm/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TimesFM init file."""
+
+print(
+    " See https://github.com/google-research/timesfm/blob/master/README.md for updated APIs."
+)
+from timesfm.timesfm_base import (
+    freq_map,
+    TimesFmCheckpoint,
+    TimesFmHparams,
+    TimesFmBase,
+)
+import sys
+
+try:
+    from timesfm.timesfm_jax import TimesFmJax as TimesFm
+    from timesfm import data_loader
+
+    print(f"Loaded Jax TimesFM, likely because python version is {sys.version}.")
+except Exception as _:
+    from timesfm.timesfm_torch import TimesFmTorch as TimesFm
+
+    print(f"Loaded PyTorch TimesFM, likely because python version is {sys.version}.")

v1/src/timesfm/data_loader.py

@@ -0,0 +1,255 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+"""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
+from sklearn.preprocessing import StandardScaler
+import tensorflow as tf
+from . import time_features
+
+
+class TimeSeriesdata(object):
+  """Data loader class."""
+
+  def __init__(
+      self,
+      data_path,
+      datetime_col,
+      num_cov_cols,
+      cat_cov_cols,
+      ts_cols,
+      train_range,
+      val_range,
+      test_range,
+      hist_len,
+      pred_len,
+      batch_size,
+      freq='H',
+      normalize=True,
+      epoch_len=None,
+      holiday=False,
+      permute=True,
+  ):
+    """Initialize objects.
+
+    Args:
+      data_path: path to csv file
+      datetime_col: column name for datetime col
+      num_cov_cols: list of numerical global covariates
+      cat_cov_cols: list of categorical global covariates
+      ts_cols: columns corresponding to ts
+      train_range: tuple of train ranges
+      val_range: tuple of validation ranges
+      test_range: tuple of test ranges
+      hist_len: historical context
+      pred_len: prediction length
+      batch_size: batch size (number of ts in a batch)
+      freq: freq of original data
+      normalize: std. normalize data or not
+      epoch_len: num iters in an epoch
+      holiday: use holiday features or not
+      permute: permute ts in train batches or not
+
+    Returns:
+      None
+    """
+    self.data_df = pd.read_csv(open(data_path, 'r'))
+    if not num_cov_cols:
+      self.data_df['ncol'] = np.zeros(self.data_df.shape[0])
+      num_cov_cols = ['ncol']
+    if not cat_cov_cols:
+      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]),
+                           inplace=True)
+    self.num_cov_cols = num_cov_cols
+    self.cat_cov_cols = cat_cov_cols
+    self.ts_cols = ts_cols
+    self.train_range = train_range
+    self.val_range = val_range
+    self.test_range = test_range
+    data_df_idx = self.data_df.index
+    date_index = data_df_idx.union(
+        pd.date_range(
+            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()
+    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.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)
+    self.normalize = normalize
+    if normalize:
+      self._normalize_data()
+    logging.info(
+        'Data Shapes: %s, %s, %s, %s',
+        self.data_mat.shape,
+        self.time_mat.shape,
+        self.num_feat_mat.shape,
+        self.cat_feat_mat.shape,
+    )
+    self.epoch_len = epoch_len
+    self.permute = permute
+
+  def _get_cat_cols(self, cat_cov_cols):
+    """Get categorical columns."""
+    cat_vars = []
+    cat_sizes = []
+    for col in cat_cov_cols:
+      dct = {x: i for i, x in enumerate(self.data_df[col].unique())}
+      cat_sizes.append(len(dct))
+      mapped = self.data_df[col].map(lambda x: dct[x]).to_numpy().transpose()  # pylint: disable=cell-var-from-loop
+      cat_vars.append(mapped)
+    return np.vstack(cat_vars), cat_sizes
+
+  def _normalize_data(self):
+    self.scaler = StandardScaler()
+    train_mat = self.data_mat[:, 0:self.train_range[1]]
+    self.scaler = self.scaler.fit(train_mat.transpose())
+    self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
+
+  def train_gen(self):
+    """Generator for training data."""
+    num_ts = len(self.ts_cols)
+    perm = np.arange(
+        self.train_range[0] + self.hist_len,
+        self.train_range[1] - self.pred_len,
+    )
+    perm = np.random.permutation(perm)
+    hist_len = self.hist_len
+    logging.info('Hist len: %s', hist_len)
+    if not self.epoch_len:
+      epoch_len = len(perm)
+    else:
+      epoch_len = self.epoch_len
+    for idx in perm[0:epoch_len]:
+      for _ in range(num_ts // self.batch_size + 1):
+        if self.permute:
+          tsidx = np.random.choice(num_ts, size=self.batch_size, replace=False)
+        else:
+          tsidx = np.arange(num_ts)
+        dtimes = np.arange(idx - hist_len, idx + self.pred_len)
+        (
+            bts_train,
+            bts_pred,
+            bfeats_train,
+            bfeats_pred,
+            bcf_train,
+            bcf_pred,
+        ) = self._get_features_and_ts(dtimes, tsidx, hist_len)
+
+        all_data = [
+            bts_train,
+            bfeats_train,
+            bcf_train,
+            bts_pred,
+            bfeats_pred,
+            bcf_pred,
+            tsidx,
+        ]
+        yield tuple(all_data)
+
+  def test_val_gen(self, mode='val', shift=1):
+    """Generator for validation/test data."""
+    if mode == 'val':
+      start = self.val_range[0]
+      end = self.val_range[1] - self.pred_len + 1
+    elif mode == 'test':
+      start = self.test_range[0]
+      end = self.test_range[1] - self.pred_len + 1
+    else:
+      raise NotImplementedError('Eval mode not implemented')
+    num_ts = len(self.ts_cols)
+    hist_len = self.hist_len
+    logging.info('Hist len: %s', hist_len)
+    perm = np.arange(start, end)
+    if self.epoch_len:
+      epoch_len = self.epoch_len
+    else:
+      epoch_len = len(perm)
+    for i in range(0, epoch_len, shift):
+      idx = perm[i]
+      for batch_idx in range(0, num_ts, self.batch_size):
+        tsidx = np.arange(batch_idx, min(batch_idx + self.batch_size, num_ts))
+        dtimes = np.arange(idx - hist_len, idx + self.pred_len)
+        (
+            bts_train,
+            bts_pred,
+            bfeats_train,
+            bfeats_pred,
+            bcf_train,
+            bcf_pred,
+        ) = self._get_features_and_ts(dtimes, tsidx, hist_len)
+        all_data = [
+            bts_train,
+            bfeats_train,
+            bcf_train,
+            bts_pred,
+            bfeats_pred,
+            bcf_pred,
+            tsidx,
+        ]
+        yield tuple(all_data)
+
+  def _get_features_and_ts(self, dtimes, tsidx, hist_len=None):
+    """Get features and ts in specified windows."""
+    if hist_len is None:
+      hist_len = self.hist_len
+    data_times = dtimes[dtimes < self.data_mat.shape[1]]
+    bdata = self.data_mat[:, data_times]
+    bts = bdata[tsidx, :]
+    bnf = self.num_feat_mat[:, data_times]
+    bcf = self.cat_feat_mat[:, data_times]
+    btf = self.time_mat[:, dtimes]
+    if bnf.shape[1] < btf.shape[1]:
+      rem_len = btf.shape[1] - bnf.shape[1]
+      rem_rep = np.repeat(bnf[:, [-1]], repeats=rem_len)
+      rem_rep_cat = np.repeat(bcf[:, [-1]], repeats=rem_len)
+      bnf = np.hstack([bnf, rem_rep.reshape(bnf.shape[0], -1)])
+      bcf = np.hstack([bcf, rem_rep_cat.reshape(bcf.shape[0], -1)])
+    bfeats = np.vstack([btf, bnf])
+    bts_train = bts[:, 0:hist_len]
+    bts_pred = bts[:, hist_len:]
+    bfeats_train = bfeats[:, 0:hist_len]
+    bfeats_pred = bfeats[:, hist_len:]
+    bcf_train = bcf[:, 0:hist_len]
+    bcf_pred = bcf[:, hist_len:]
+    return bts_train, bts_pred, bfeats_train, bfeats_pred, bcf_train, bcf_pred
+
+  def tf_dataset(self, mode='train', shift=1):
+    """Tensorflow Dataset."""
+    if mode == 'train':
+      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 +
+                         [tf.int32] * 2)
+    dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
+    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+    return dataset

v1/src/timesfm/patched_decoder.py

@@ -0,0 +1,543 @@
+# 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.
+"""Pax ML model for patched time-series decoder.
+
+The file implements Residual MLPs, Patched Decoder layers and PAX ML models.
+"""
+
+import dataclasses
+from typing import Optional, Tuple
+
+import einshape as es
+from jax import lax
+import jax.numpy as jnp
+from praxis import base_layer
+from praxis import base_model
+from praxis import layers
+from praxis import pax_fiddle
+from praxis import py_utils
+from praxis import pytypes
+from praxis.layers import activations
+from praxis.layers import embedding_softmax
+from praxis.layers import linears
+from praxis.layers import normalizations
+from praxis.layers import stochastics
+from praxis.layers import transformers
+
+# PAX shortcuts
+NestedMap = py_utils.NestedMap
+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]
+
+# NestedMap keys
+_INPUT_TS = "input_ts"
+_TARGET_FUTURE = "actual_ts"
+_INPUT_PADDING = "input_padding"
+_OUTPUT_TS = "output_ts"
+_FREQ = "freq"
+_OUTPUT_TOKENS = "output_tokens"
+_STATS = "stats"
+
+# Small numerical value.
+_TOLERANCE = 1e-7
+
+
+def _shift_padded_seq(mask: JTensor, seq: JTensor) -> JTensor:
+  """Shifts rows of seq based on the first 0 in each row of the mask."""
+  num = seq.shape[1]
+
+  # Find the index of the first 0 in each row of the mask
+  first_zero_idx = jnp.argmin(mask, axis=1)
+
+  # Create a range array for indexing
+  idx_range = jnp.arange(num)
+
+  def shift_row(carry, x):
+    seq_row, shift = x
+    shifted_idx = (idx_range - shift) % num
+    shifted_row = seq_row[shifted_idx]
+    return carry, shifted_row
+
+  # Use lax.scan to shift each row of seq based on the corresponding
+  # first_zero_idx.
+  _, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx))
+
+  return shifted_seq
+
+
+class ResidualBlock(base_layer.BaseLayer):
+  """Simple feedforward block with residual connection.
+
+  Attributes:
+    input_dims: input dimension.
+    hidden_dims: hidden dimension.
+    output_dims: output dimension.
+    dropout_prob: dropout probability.
+    layer_norm: whether to use layer norm or not.
+    dropout_tpl: config for dropout.
+    ln_tpl: config for layer norm.
+    act_tpl: config for activation in hidden layer.
+  """
+
+  input_dims: int = 0
+  hidden_dims: int = 0
+  output_dims: int = 0
+  dropout_prob: float = 0.0
+  layer_norm: bool = False
+  dropout_tpl: LayerTpl = template_field(stochastics.Dropout)
+  ln_tpl: LayerTpl = template_field(normalizations.LayerNorm)
+  act_tpl: LayerTpl = template_field(activations.Swish)
+
+  def setup(self):
+    lnorm_tpl = self.ln_tpl.clone()
+    lnorm_tpl.dim = self.output_dims
+    self.create_child("ln_layer", lnorm_tpl)
+
+    dropout_tpl = self.dropout_tpl.clone()
+    dropout_tpl.keep_prob = 1.0 - self.dropout_prob
+    self.create_child("dropout", dropout_tpl)
+
+    self.create_child(
+        "hidden_layer",
+        pax_fiddle.Config(
+            linears.FeedForward,
+            input_dims=self.input_dims,
+            output_dims=self.hidden_dims,
+            activation_tpl=self.act_tpl.clone(),
+        ),
+    )
+
+    self.create_child(
+        "output_layer",
+        pax_fiddle.Config(
+            linears.FeedForward,
+            input_dims=self.hidden_dims,
+            output_dims=self.output_dims,
+            activation_tpl=pax_fiddle.Config(activations.Identity),
+        ),
+    )
+
+    self.create_child(
+        "residual_layer",
+        pax_fiddle.Config(
+            linears.FeedForward,
+            input_dims=self.input_dims,
+            output_dims=self.output_dims,
+            activation_tpl=pax_fiddle.Config(activations.Identity),
+        ),
+    )
+
+  def __call__(self, inputs: JTensor) -> JTensor:
+    hidden = self.hidden_layer(inputs)
+    output = self.output_layer(hidden)
+    output = self.dropout(output)
+    residual = self.residual_layer(inputs)
+    if self.layer_norm:
+      return self.ln_layer(output + residual)
+    else:
+      return output + residual
+
+
+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.
+
+  Args:
+    inputs: A JAX array of shape [b, n, p].
+    padding: A JAX array of shape [b, n, p] with values 0 or 1.
+
+  Returns:
+    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.
+  """
+  # Selecting the first pad with more than 3 unpadded values.
+  pad_sum = jnp.sum(1 - padding, axis=2)
+
+  def _get_patch_index(arr: JTensor):
+    indices = jnp.argmax(arr >= 3, axis=1)
+    row_sum = (arr >= 3).sum(axis=1)
+    return jnp.where(row_sum == 0, arr.shape[1] - 1, indices)
+
+  patch_indices = _get_patch_index(pad_sum)
+  bidxs = jnp.arange(inputs.shape[0])
+
+  arr = inputs[bidxs, patch_indices, :]
+  pad = padding[bidxs, patch_indices, :]
+
+  # Create a mask where P is 0
+  mask = 1 - pad
+
+  # Calculate the number of valid elements
+  num_valid_elements = jnp.sum(mask, axis=1)
+
+  num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements)
+
+  # 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)
+
+  # Calculate the masked mean and standard deviation
+  masked_mean = masked_sum / num_valid_elements
+  masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
+  masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
+  masked_std = jnp.sqrt(masked_var)
+
+  return masked_mean, masked_std
+
+
+def _create_quantiles() -> list[float]:
+  """Returns the quantiles for forecasting."""
+  return DEFAULT_QUANTILES
+
+
+class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
+  """Patch decoder layer for time-series foundation model.
+
+  Attributes:
+    patch_len: length of input patches.
+    horizon_len: length of output patches. Referred to as `output_patch_len`
+      during inference.
+    model_dims: model dimension of stacked transformer layer.
+    hidden_dims: hidden dimensions in fully connected layers.
+    quantiles: list of quantiles for non prob model.
+    residual_block_tpl: config for residual block.
+    stacked_transformer_params_tpl: config for stacked transformer.
+    use_freq: whether to use frequency encoding.
+
+  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
+  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.
+  """
+
+  patch_len: int = 0
+  horizon_len: int = 0
+  model_dims: int = 0
+  hidden_dims: int = 0
+  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)
+  use_freq: bool = True
+  use_pos_emb: bool = True
+
+  def setup(self) -> None:
+    """Construct the model."""
+    num_outputs = len(self.quantiles) + 1
+
+    stl = self.stacked_transformer_params_tpl.clone()
+    stl.model_dims = self.model_dims
+    stl.hidden_dims = self.hidden_dims
+    stl.mask_self_attention = True
+
+    self.create_child("stacked_transformer_layer", stl)
+
+    input_resl = self.residual_block_tpl.clone()
+    ff_in_dims = 2 * self.patch_len
+    input_resl.input_dims = ff_in_dims
+    input_resl.hidden_dims = self.hidden_dims
+    input_resl.output_dims = self.model_dims
+    self.create_child(
+        "input_ff_layer",
+        input_resl,
+    )
+
+    horizon_resl = self.residual_block_tpl.clone()
+    horizon_resl.input_dims = self.model_dims
+    horizon_resl.hidden_dims = self.hidden_dims
+    horizon_resl.output_dims = self.horizon_len * num_outputs
+    self.create_child(
+        "horizon_ff_layer",
+        horizon_resl,
+    )
+
+    self.create_child(
+        "position_emb",
+        pax_fiddle.Config(layers.PositionalEmbedding,
+                          embedding_dims=self.model_dims),
+    )
+
+    if self.use_freq:
+      self.create_child(
+          "freq_emb",
+          pax_fiddle.Config(
+              embedding_softmax.Embedding,
+              num_classes=3,
+              input_dims=self.model_dims,
+          ),
+      )
+
+  def transform_decode_state(
+      self, transform_fn: base_layer.DecodeStateTransformFn) -> None:
+    """Transforms all decode state variables based on transform_fn."""
+    self.stacked_transformer_layer.transform_decode_state(transform_fn)
+
+  def _forward_transform(
+      self, inputs: JTensor,
+      patched_pads: JTensor) -> Tuple[JTensor, Tuple[JTensor, JTensor]]:
+    """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)
+    return outputs, (mu, sigma)
+
+  def _reverse_transform(self, outputs: JTensor,
+                         stats: Tuple[JTensor, JTensor]) -> JTensor:
+    """Output is of shape [B, N, P, Q]."""
+    mu, sigma = stats
+    return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
+
+  def _preprocess_input(
+      self,
+      input_ts: JTensor,
+      input_padding: JTensor,
+      pos_emb: Optional[JTensor] = None,
+  ) -> Tuple[JTensor, JTensor, Optional[Tuple[JTensor, JTensor]], JTensor]:
+    """Preprocess input for stacked transformer."""
+    # Reshape into patches.
+    patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len)
+    patched_pads = es.jax_einshape("b(np)->bnp",
+                                   input_padding,
+                                   p=self.patch_len)
+    patched_inputs = jnp.where(
+        jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs)
+    patched_pads = jnp.where(
+        jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads)
+    patched_inputs, stats = self._forward_transform(patched_inputs,
+                                                    patched_pads)
+
+    # B x N x D
+    patched_inputs = patched_inputs * (1.0 - patched_pads)
+    concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1)
+    model_input = self.input_ff_layer(concat_inputs)
+    # A patch should not be padded even if there is at least one zero.
+    patched_padding = jnp.min(patched_pads, axis=-1)
+
+    if self.use_pos_emb:
+      if pos_emb is None:
+        position_emb = self.position_emb(seq_length=model_input.shape[1])
+      else:
+        position_emb = pos_emb
+      if self.do_eval:
+        if position_emb.shape[0] != model_input.shape[0]:
+          position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0)
+        position_emb = _shift_padded_seq(patched_padding, position_emb)
+      model_input += position_emb
+
+    return model_input, patched_padding, stats, patched_inputs
+
+  def _postprocess_output(
+      self,
+      model_output: JTensor,
+      num_outputs: int,
+      stats: Tuple[JTensor, JTensor],
+  ) -> JTensor:
+    """Postprocess output of stacked transformer."""
+    # B x N x (H.Q)
+    output_ts = self.horizon_ff_layer(model_output)
+    output_ts = es.jax_einshape("bn(hq)->bnhq",
+                                output_ts,
+                                q=num_outputs,
+                                h=self.horizon_len)
+    return self._reverse_transform(output_ts, stats)
+
+  def __call__(self, inputs: NestedMap) -> NestedMap:
+    """PatchTST call.
+
+    Args:
+      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
+        contains padding map.
+
+    Returns:
+      A nested map with two keys:
+      (1) 'output_tokens' of shape [B, N, D].
+      (2) 'output_ts' of shape [B, N, H, Q]
+      (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(
+        input_ts=input_ts,
+        input_padding=input_padding,
+    )
+    if self.use_freq:
+      freq = inputs[_FREQ].astype(jnp.int32)
+      f_emb = self.freq_emb(freq)  # B x 1 x D
+      f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1)
+      model_input += f_emb
+    model_output = self.stacked_transformer_layer(model_input, patched_padding)
+
+    output_ts = self._postprocess_output(model_output, num_outputs, stats)
+    return NestedMap({
+        _OUTPUT_TOKENS: model_output,
+        _OUTPUT_TS: output_ts,
+        _STATS: stats
+    })
+
+  def decode(
+      self,
+      inputs: NestedMap,
+      horizon_len: int,
+      output_patch_len: Optional[int] = None,
+      max_len: int | None = None,
+      return_forecast_on_context: bool = False,
+  ) -> tuple[JTensor, JTensor]:
+    """Auto-regressive decoding without caching.
+
+    Args:
+      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.
+      horizon_len: prediction length.
+      output_patch_len: output length to be fetched from one step of
+        auto-regressive decoding.
+      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:
+      Tuple of two forecasting results:
+      - Point (mean) output predictions as a tensor with shape B x H'.
+      - Full predictions (mean and quantiles) as a tensor with shape
+        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]
+    context_len = final_out.shape[1]
+    paddings = inputs[_INPUT_PADDING]
+    if max_len is None:
+      max_len = context_len
+    if self.use_freq:
+      freq = inputs[_FREQ].astype(jnp.int32)
+    else:
+      freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32)
+    full_outputs = []
+    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}")
+    if output_patch_len is None:
+      output_patch_len = self.horizon_len
+    num_decode_patches = (horizon_len + output_patch_len -
+                          1) // output_patch_len
+    for step_index in range(num_decode_patches):
+      current_padding = paddings[:, 0:final_out.shape[1]]
+      input_ts = final_out[:, -max_len:]
+      input_padding = current_padding[:, -max_len:]
+      model_input = NestedMap(
+          input_ts=input_ts,
+          input_padding=input_padding,
+          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(new_full_ts)
+      final_out = jnp.concatenate([final_out, new_ts], axis=-1)
+
+    if return_forecast_on_context:
+      # `full_outputs` indexing starts at after the first input patch.
+      full_outputs = jnp.concatenate(full_outputs,
+                                     axis=1)[:, :(context_len - self.patch_len +
+                                                  horizon_len), :]
+    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:
+    core_layer_tpl: config for core layer.
+    freq: freq to finetune on.
+  """
+
+  core_layer_tpl: LayerTpl = template_field(PatchedTimeSeriesDecoder)
+  freq: int = 0
+
+  def setup(self) -> None:
+    self.create_child("core_layer", self.core_layer_tpl)
+
+  def compute_predictions(self, input_batch: NestedMap) -> NestedMap:
+    input_ts = input_batch[_INPUT_TS]
+    input_padding = jnp.zeros_like(input_ts)
+    context_len = input_ts.shape[1]
+    input_patch_len = self.core_layer_tpl.patch_len
+    context_pad = ((context_len + input_patch_len - 1) //
+                   input_patch_len) * input_patch_len - context_len
+
+    input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
+    input_padding = jnp.pad(input_padding, [(0, 0), (context_pad, 0)],
+                            constant_values=1)
+    freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
+    new_input_batch = NestedMap(
+        input_ts=input_ts,
+        input_padding=input_padding,
+        freq=freq,
+    )
+    return self.core_layer(new_input_batch)
+
+  def _quantile_loss(self, pred: JTensor, actual: JTensor,
+                     quantile: float) -> JTensor:
+    """Calculates quantile loss.
+
+    Args:
+      pred: B x T
+      actual: B x T
+      quantile: quantile at which loss is computed.
+
+    Returns:
+      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,
+                   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], :]
+    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)
+    loss = loss.mean()
+    loss_weight = jnp.array(1.0, dtype=jnp.float32)
+    per_example_out = NestedMap()
+    return {"avg_qloss": (loss, loss_weight)}, per_example_out

v1/src/timesfm/pytorch_patched_decoder.py

@@ -0,0 +1,801 @@
+# 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.
+"""Pytorch version of patched decoder."""
+
+import dataclasses
+import math
+from typing import List, Tuple
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+def create_quantiles() -> list[float]:
+  return [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+
+
+@dataclasses.dataclass
+class TimesFMConfig:
+  """Config for initializing timesfm patched_decoder class."""
+
+  # The number of blocks in the model.
+  num_layers: int = 20
+  # The number of attention heads used in the attention layers of the model.
+  num_heads: int = 16
+  # The number of key-value heads for implementing attention.
+  num_kv_heads: int = 16
+  # The hidden size of the model.
+  hidden_size: int = 1280
+  # The dimension of the MLP representations.
+  intermediate_size: int = 1280
+  # The number of head dimensions.
+  head_dim: int = 80
+  # The epsilon used by the rms normalization layers.
+  rms_norm_eps: float = 1e-6
+  # Patch length
+  patch_len: int = 32
+  # Horizon length
+  horizon_len: int = 128
+  # quantiles
+  quantiles: List[float] = dataclasses.field(default_factory=create_quantiles)
+  # Padding value
+  pad_val: float = 1123581321.0
+  # Tolerance
+  tolerance: float = 1e-6
+  # The dtype of the weights.
+  dtype: str = "bfloat32"
+  # use positional embedding
+  use_positional_embedding: bool = True
+
+
+def _masked_mean_std(
+    inputs: torch.Tensor,
+    padding: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+  """Calculates mean and standard deviation of `inputs` across axis 1.
+
+  It excludes values where `padding` is 1.
+
+  Args:
+    inputs: A PyTorch tensor of shape [b, n, p].
+    padding: A PyTorch tensor of shape [b, n, p] with values 0 or 1.
+
+  Returns:
+    A tuple containing the mean and standard deviation.
+    We return the statistics of the first patch with more than three non-padded
+    values.
+  """
+  # Selecting the first patch with more than 3 unpadded values.
+  pad_sum = torch.sum(1 - padding, dim=2)
+
+  def _get_patch_index(arr: torch.Tensor):
+    indices = torch.argmax((arr >= 3).to(torch.int32), dim=1)
+    row_sum = (arr >= 3).to(torch.int32).sum(dim=1)
+    return torch.where(row_sum == 0, arr.shape[1] - 1, indices)
+
+  patch_indices = _get_patch_index(pad_sum)
+  bidxs = torch.arange(inputs.shape[0])
+
+  arr = inputs[bidxs, patch_indices, :]
+  pad = padding[bidxs, patch_indices, :]
+
+  # Create a mask where padding is 0
+  mask = 1 - pad
+
+  # Calculate the number of valid elements
+  num_valid_elements = torch.sum(mask, dim=1)
+  num_valid_elements = torch.where(
+      num_valid_elements == 0,
+      torch.tensor(1,
+                   dtype=num_valid_elements.dtype,
+                   device=num_valid_elements.device),
+      num_valid_elements,
+  )
+
+  # Calculate the masked sum and squared sum
+  masked_sum = torch.sum(arr * mask, dim=1)
+  masked_squared_sum = torch.sum((arr * mask)**2, dim=1)
+
+  # Calculate the masked mean and standard deviation
+  masked_mean = masked_sum / num_valid_elements
+  masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
+  masked_var = torch.where(
+      masked_var < 0.0,
+      torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),
+      masked_var,
+  )
+  masked_std = torch.sqrt(masked_var)
+
+  return masked_mean, masked_std
+
+
+def _shift_padded_seq(mask: torch.Tensor, seq: torch.Tensor) -> torch.Tensor:
+  """Shifts rows of seq based on the first 0 in each row of the mask.
+
+  Args:
+    mask: mask tensor of shape [B, N]
+    seq: seq tensor of shape [B, N, P]
+
+  Returns:
+    Returns the shifted sequence.
+  """
+  batch_size, num_seq, feature_dim = seq.shape
+
+  new_mask: torch.BoolTensor = mask == 0
+
+  # Use argmax to find the first True value in each row
+  indices = new_mask.to(torch.int32).argmax(dim=1)
+
+  # Handle rows with all zeros
+  indices[~new_mask.any(dim=1)] = -1
+
+  # Create index ranges for each sequence in the batch
+  idx_range = (torch.arange(num_seq).to(
+      seq.device).unsqueeze(0).unsqueeze(-1).expand(batch_size, -1,
+                                                    feature_dim))
+
+  # Calculate shifted indices for each element in each sequence
+  shifted_idx = (idx_range - indices[:, None, None]) % num_seq
+
+  # Gather values from seq using shifted indices
+  shifted_seq = seq.gather(1, shifted_idx)
+
+  return shifted_seq
+
+
+def get_large_negative_number(dtype: torch.dtype) -> torch.Tensor:
+  """Returns a large negative value for the given dtype."""
+  if dtype.is_floating_point:
+    dtype_max = torch.finfo(dtype).max
+  else:
+    dtype_max = torch.iinfo(dtype).max
+  return torch.tensor(-0.7 * dtype_max, dtype=dtype)
+
+
+def apply_mask_to_logits(logits: torch.Tensor,
+                         mask: torch.Tensor) -> torch.Tensor:
+  """Applies a floating-point mask to a set of logits.
+
+  Args:
+      logits: A torch.Tensor of logit values.
+      mask: A torch.Tensor (float32) of mask values with the encoding described
+        in the function documentation.
+
+  Returns:
+      Masked logits.
+  """
+
+  min_value = get_large_negative_number(logits.dtype)
+
+  return torch.where((mask >= min_value * 0.5), logits, min_value)
+
+
+def convert_paddings_to_mask(
+    paddings: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
+  """Converts binary paddings to a logit mask ready to add to attention matrix.
+
+  Args:
+      paddings: binary torch.Tensor of shape [B, T], with 1 denoting padding
+        token.
+      dtype: data type of the input.
+
+  Returns:
+      A torch.Tensor of shape [B, 1, 1, T] ready to add to attention logits.
+  """
+  attention_mask = paddings.detach().clone()
+  attention_mask = attention_mask[:, None, None, :]  # Equivalent to jnp.newaxis
+  attention_mask *= get_large_negative_number(dtype)
+  return attention_mask
+
+
+def causal_mask(input_t: torch.Tensor) -> torch.Tensor:
+  """Computes and returns causal mask.
+
+  Args:
+      input_t: A torch.Tensor of shape [B, T, D].
+
+  Returns:
+      An attention_mask torch.Tensor of shape [1, 1, T, T]. Attention mask has
+      already been converted to large negative values.
+  """
+  assert input_t.dtype.is_floating_point, input_t.dtype
+  large_negative_number = get_large_negative_number(input_t.dtype)
+  t = input_t.shape[1]
+  col_idx = torch.arange(t).unsqueeze(0).repeat(t, 1)
+  row_idx = torch.arange(t).unsqueeze(1).repeat(1, t)
+  mask = (row_idx < col_idx).to(input_t.dtype) * large_negative_number
+  return (mask.unsqueeze(0).unsqueeze(0).to(input_t.device)
+         )  # Equivalent to jnp.newaxis
+
+
+def merge_masks(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+  """Merges 2 masks.
+
+  logscale mask is expected but 0/1 mask is also fine.
+
+  Args:
+      a: torch.Tensor of shape [1|B, 1, 1|T, S].
+      b: torch.Tensor of shape [1|B, 1, 1|T, S].
+
+  Returns:
+      torch.Tensor of shape [1|B, 1, 1|T, S].
+  """
+
+  def expand_t(key_mask):
+    query_mask = key_mask.transpose(-1, -2)  # Equivalent of jnp.transpose
+    return torch.minimum(query_mask, key_mask)
+
+  if a.shape[2] != b.shape[2]:
+    if a.shape[2] == 1:
+      a = expand_t(a)
+    else:
+      assert b.shape[2] == 1
+      b = expand_t(b)
+
+  assert a.shape[1:] == b.shape[1:], f"a.shape={a.shape}, b.shape={b.shape}."
+  return torch.minimum(a, b)  # Element-wise minimum, similar to jnp.minimum
+
+
+class ResidualBlock(nn.Module):
+  """TimesFM residual block."""
+
+  def __init__(
+      self,
+      input_dims,
+      hidden_dims,
+      output_dims,
+  ):
+    super(ResidualBlock, self).__init__()
+    self.input_dims = input_dims
+    self.hidden_dims = hidden_dims
+    self.output_dims = output_dims
+
+    # Hidden Layer
+    self.hidden_layer = nn.Sequential(
+        nn.Linear(input_dims, hidden_dims),
+        nn.SiLU(),
+    )
+
+    # Output Layer
+    self.output_layer = nn.Linear(hidden_dims, output_dims)
+    # Residual Layer
+    self.residual_layer = nn.Linear(input_dims, output_dims)
+
+  def forward(self, x):
+    hidden = self.hidden_layer(x)
+    output = self.output_layer(hidden)
+    residual = self.residual_layer(x)
+    return output + residual
+
+
+class RMSNorm(torch.nn.Module):
+  """Pax rms norm in pytorch."""
+
+  def __init__(
+      self,
+      dim: int,
+      eps: float = 1e-6,
+      add_unit_offset: bool = False,
+  ):
+    super().__init__()
+    self.eps = eps
+    self.add_unit_offset = add_unit_offset
+    self.weight = nn.Parameter(torch.zeros(dim))
+
+  def _norm(self, x):
+    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+  def forward(self, x):
+    output = self._norm(x.float())
+    if self.add_unit_offset:
+      output = output * (1 + self.weight.float())
+    else:
+      output = output * self.weight.float()
+    return output.type_as(x)
+
+
+class TransformerMLP(nn.Module):
+  """Pax transformer MLP in pytorch."""
+
+  def __init__(
+      self,
+      hidden_size: int,
+      intermediate_size: int,
+  ):
+    super().__init__()
+    self.gate_proj = nn.Linear(hidden_size, intermediate_size)
+    self.down_proj = nn.Linear(intermediate_size, hidden_size)
+    self.layer_norm = nn.LayerNorm(normalized_shape=hidden_size, eps=1e-6)
+
+  def forward(self, x, paddings=None):
+    gate_inp = self.layer_norm(x)
+    gate = self.gate_proj(gate_inp)
+    gate = F.relu(gate)
+    outputs = self.down_proj(gate)
+    if paddings is not None:
+      outputs = outputs * (1.0 - paddings[:, :, None])
+    return outputs + x
+
+
+class TimesFMAttention(nn.Module):
+  """Implements the attention used in TimesFM."""
+
+  def __init__(
+      self,
+      hidden_size: int,
+      num_heads: int,
+      num_kv_heads: int,
+      head_dim: int,
+  ):
+    super().__init__()
+
+    self.num_heads = num_heads
+    self.num_kv_heads = num_kv_heads
+
+    assert self.num_heads % self.num_kv_heads == 0
+    self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+    self.hidden_size = hidden_size
+    self.head_dim = head_dim
+
+    self.q_size = self.num_heads * self.head_dim
+    self.kv_size = self.num_kv_heads * self.head_dim
+    self.scaling = nn.Parameter(
+        torch.empty((self.head_dim,), dtype=torch.float32),)
+
+    self.qkv_proj = nn.Linear(
+        self.hidden_size,
+        (self.num_heads + 2 * self.num_kv_heads) * self.head_dim,
+    )
+    self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size)
+
+  def _per_dim_scaling(self, query: torch.Tensor) -> torch.Tensor:
+    # [batch_size, n_local_heads, input_len, head_dim]
+    r_softplus_0 = 1.442695041
+    softplus_func = torch.nn.Softplus()
+    scale = r_softplus_0 / math.sqrt(self.head_dim)
+    scale = scale * softplus_func(self.scaling)
+    return query * scale[None, None, None, :]
+
+  def forward(
+      self,
+      hidden_states: torch.Tensor,
+      mask: torch.Tensor,
+      kv_write_indices: torch.Tensor | None = None,
+      kv_cache: Tuple[torch.Tensor, torch.Tensor] | None = None,
+  ) -> torch.Tensor:
+    hidden_states_shape = hidden_states.shape
+    assert len(hidden_states_shape) == 3
+
+    batch_size, input_len, _ = hidden_states_shape
+
+    qkv = self.qkv_proj(hidden_states)
+    xq, xk, xv = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+    xq = xq.view(batch_size, -1, self.num_heads, self.head_dim)
+    xk = xk.view(batch_size, -1, self.num_kv_heads, self.head_dim)
+    xv = xv.view(batch_size, -1, self.num_kv_heads, self.head_dim)
+    xq = self._per_dim_scaling(xq)
+
+    # Write new kv cache.
+    # [batch_size, input_len, n_local_kv_heads, head_dim]
+    if kv_cache is not None and kv_write_indices is not None:
+      k_cache, v_cache = kv_cache
+      k_cache.index_copy_(1, kv_write_indices, xk)
+      v_cache.index_copy_(1, kv_write_indices, xv)
+
+      key = k_cache
+      value = v_cache
+    else:
+      key = xk
+      value = xv
+    if self.num_kv_heads != self.num_heads:
+      # [batch_size, max_seq_len, n_local_heads, head_dim]
+      key = torch.repeat_interleave(key, self.num_queries_per_kv, dim=2)
+      value = torch.repeat_interleave(value, self.num_queries_per_kv, dim=2)
+
+    # [batch_size, n_local_heads, input_len, head_dim]
+    q = xq.transpose(1, 2)
+    # [batch_size, n_local_heads, max_seq_len, head_dim]
+    k = key.transpose(1, 2)
+    v = value.transpose(1, 2)
+
+    # [batch_size, n_local_heads, input_len, max_seq_len]
+    scores = torch.matmul(q, k.transpose(2, 3))
+    scores = scores + mask
+    scores = F.softmax(scores.float(), dim=-1).type_as(q)
+
+    # [batch_size, n_local_heads, input_len, head_dim]
+    output = torch.matmul(scores, v)
+    # return scores, output.transpose(1, 2).contiguous()
+
+    # [batch_size, input_len, hidden_dim]
+    output = output.transpose(1, 2).contiguous().view(batch_size, input_len, -1)
+    output = self.o_proj(output)
+    return scores, output
+
+
+class TimesFMDecoderLayer(nn.Module):
+  """Transformer layer."""
+
+  def __init__(
+      self,
+      hidden_size: int,
+      intermediate_size: int,
+      num_heads: int,
+      num_kv_heads: int,
+      head_dim: int,
+      rms_norm_eps: float = 1e-6,
+  ):
+    super().__init__()
+    self.self_attn = TimesFMAttention(
+        hidden_size=hidden_size,
+        num_heads=num_heads,
+        num_kv_heads=num_kv_heads,
+        head_dim=head_dim,
+    )
+    self.mlp = TransformerMLP(
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+    )
+    self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+
+  def forward(
+      self,
+      hidden_states: torch.Tensor,
+      mask: torch.Tensor,
+      paddings: torch.Tensor,
+      kv_write_indices: torch.Tensor | None = None,
+      kv_cache: Tuple[torch.Tensor, torch.Tensor] | None = None,
+  ) -> torch.Tensor:
+    # Self Attention
+    residual = hidden_states
+    hidden_states = self.input_layernorm(hidden_states)
+    scores, hidden_states = self.self_attn(
+        hidden_states=hidden_states,
+        mask=mask,
+        kv_write_indices=kv_write_indices,
+        kv_cache=kv_cache,
+    )
+    hidden_states = residual + hidden_states
+
+    # MLP
+    hidden_states = self.mlp(hidden_states, paddings=paddings)
+
+    return scores, hidden_states
+
+
+class StackedDecoder(nn.Module):
+  """Stacked transformer layer."""
+
+  def __init__(
+      self,
+      hidden_size: int,
+      intermediate_size: int,
+      num_heads: int,
+      num_kv_heads: int,
+      head_dim: int,
+      num_layers: int,
+      rms_norm_eps: float = 1e-6,
+  ):
+    super().__init__()
+
+    self.layers = nn.ModuleList()
+    for _ in range(num_layers):
+      self.layers.append(
+          TimesFMDecoderLayer(
+              hidden_size=hidden_size,
+              intermediate_size=intermediate_size,
+              num_heads=num_heads,
+              num_kv_heads=num_kv_heads,
+              head_dim=head_dim,
+              rms_norm_eps=rms_norm_eps,
+          ))
+
+  def forward(
+      self,
+      hidden_states: torch.Tensor,
+      paddings: torch.Tensor,
+      kv_write_indices: torch.Tensor | None = None,
+      kv_caches: List[Tuple[torch.Tensor, torch.Tensor]] | None = None,
+  ) -> torch.Tensor:
+    padding_mask = convert_paddings_to_mask(paddings, hidden_states.dtype)
+    atten_mask = causal_mask(hidden_states)
+    mask = merge_masks(padding_mask, atten_mask)
+    for i in range(len(self.layers)):
+      layer = self.layers[i]
+      kv_cache = kv_caches[i] if kv_caches is not None else None
+      _, hidden_states = layer(
+          hidden_states=hidden_states,
+          mask=mask,
+          paddings=paddings,
+          kv_write_indices=kv_write_indices,
+          kv_cache=kv_cache,
+      )
+    return hidden_states
+
+
+class PositionalEmbedding(torch.nn.Module):
+  """Generates position embedding for a given 1-d sequence.
+
+  Attributes:
+      min_timescale: Start of the geometric index. Determines the periodicity of
+        the added signal.
+      max_timescale: End of the geometric index. Determines the frequency of the
+        added signal.
+      embedding_dims: Dimension of the embedding to be generated.
+  """
+
+  def __init__(
+      self,
+      embedding_dims: int,
+      min_timescale: int = 1,
+      max_timescale: int = 10_000,
+  ) -> None:
+    super().__init__()
+    self.min_timescale = min_timescale
+    self.max_timescale = max_timescale
+    self.embedding_dims = embedding_dims
+
+  def forward(self, seq_length=None, position=None):
+    """Generates a Tensor of sinusoids with different frequencies.
+
+    Args:
+        seq_length: an optional Python int defining the output sequence length.
+          if the `position` argument is specified.
+        position:   [B, seq_length], optional position for each token in the
+          sequence, only required when the sequence is packed.
+
+    Returns:
+        [B, seqlen, D] if `position` is specified, else [1, seqlen, D]
+    """
+    if position is None:
+      assert seq_length is not None
+      # [1, seqlen]
+      position = torch.arange(seq_length, dtype=torch.float32).unsqueeze(0)
+    else:
+      assert position.ndim == 2, position.shape
+
+    num_timescales = self.embedding_dims // 2
+    log_timescale_increment = math.log(
+        float(self.max_timescale) / float(self.min_timescale)) / max(
+            num_timescales - 1, 1)
+    inv_timescales = self.min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float32) *
+        -log_timescale_increment)
+    scaled_time = position.unsqueeze(2) * inv_timescales.unsqueeze(0).unsqueeze(
+        0)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=2)
+    # Padding to ensure correct embedding dimension
+    signal = F.pad(signal, (0, 0, 0, self.embedding_dims % 2))
+    return signal
+
+
+class PatchedTimeSeriesDecoder(nn.Module):
+  """Patched time-series decoder."""
+
+  def __init__(self, config: TimesFMConfig):
+    super().__init__()
+    self.config = config
+    self.input_ff_layer = ResidualBlock(
+        input_dims=2 * config.patch_len,
+        output_dims=config.hidden_size,
+        hidden_dims=config.intermediate_size,
+    )
+    self.freq_emb = nn.Embedding(num_embeddings=3,
+                                 embedding_dim=config.hidden_size)
+    self.horizon_ff_layer = ResidualBlock(
+        input_dims=config.hidden_size,
+        output_dims=config.horizon_len * (1 + len(config.quantiles)),
+        hidden_dims=config.intermediate_size,
+    )
+    self.stacked_transformer = StackedDecoder(
+        hidden_size=self.config.hidden_size,
+        intermediate_size=self.config.intermediate_size,
+        num_heads=self.config.num_heads,
+        num_kv_heads=self.config.num_kv_heads,
+        head_dim=self.config.head_dim,
+        num_layers=self.config.num_layers,
+        rms_norm_eps=self.config.rms_norm_eps,
+    )
+    if self.config.use_positional_embedding:
+      self.position_emb = PositionalEmbedding(self.config.hidden_size)
+
+  def _forward_transform(
+      self, inputs: torch.Tensor, patched_pads: torch.Tensor
+  ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+    """Input is of shape [B, N, P]."""
+    mu, sigma = _masked_mean_std(inputs, patched_pads)
+    sigma = torch.where(
+        sigma < self.config.tolerance,
+        torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),
+        sigma,
+    )
+
+    # Normalize each patch
+    outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
+    outputs = torch.where(
+        torch.abs(inputs - self.config.pad_val) < self.config.tolerance,
+        torch.tensor(self.config.pad_val,
+                     dtype=outputs.dtype,
+                     device=outputs.device),
+        outputs,
+    )
+    return outputs, (mu, sigma)
+
+  def _reverse_transform(
+      self, outputs: torch.Tensor, stats: tuple[torch.Tensor,
+                                                torch.Tensor]) -> torch.Tensor:
+    """Output is of shape [B, N, P, Q]."""
+    mu, sigma = stats
+    return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
+
+  def _preprocess_input(
+      self,
+      input_ts: torch.Tensor,
+      input_padding: torch.Tensor,
+  ) -> tuple[
+      torch.Tensor,
+      torch.Tensor,
+      tuple[torch.Tensor, torch.Tensor] | None,
+      torch.Tensor,
+  ]:
+    """Preprocess input for stacked transformer."""
+
+    # Reshape into patches (using view for efficiency)
+    bsize = input_ts.shape[0]
+    patched_inputs = input_ts.view(bsize, -1, self.config.patch_len)
+    patched_pads = input_padding.view(bsize, -1, self.config.patch_len)
+
+    patched_inputs = torch.where(
+        torch.abs(patched_pads - 1.0) < self.config.tolerance,
+        torch.tensor(0.0,
+                     dtype=patched_inputs.dtype,
+                     device=patched_inputs.device),
+        patched_inputs,
+    )
+    patched_pads = torch.where(
+        torch.abs(patched_inputs - self.config.pad_val) < self.config.tolerance,
+        torch.tensor(1.0, dtype=patched_pads.dtype, device=patched_pads.device),
+        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 = torch.cat([patched_inputs, patched_pads], dim=-1)
+    model_input = self.input_ff_layer(concat_inputs)
+
+    # A patch should not be padded even if there is at least one zero.
+    patched_padding = torch.min(patched_pads,
+                                dim=-1)[0]  # Get the values from the min result
+    if self.config.use_positional_embedding:
+      pos_emb = self.position_emb(model_input.shape[1]).to(model_input.device)
+      pos_emb = torch.concat([pos_emb] * model_input.shape[0], dim=0)
+      pos_emb = _shift_padded_seq(patched_padding, pos_emb)
+      model_input += pos_emb
+
+    return model_input, patched_padding, stats, patched_inputs
+
+  def _postprocess_output(
+      self,
+      model_output: torch.Tensor,
+      num_outputs: int,
+      stats: tuple[torch.Tensor, torch.Tensor],
+  ) -> torch.Tensor:
+    """Postprocess output of stacked transformer."""
+
+    # B x N x (H.Q)
+    output_ts = self.horizon_ff_layer(model_output)
+
+    # Reshape using view
+    b, n, _ = output_ts.shape
+    output_ts = output_ts.view(b, n, self.config.horizon_len, num_outputs)
+
+    return self._reverse_transform(output_ts, stats)
+
+  def forward(
+      self,
+      input_ts: torch.Tensor,
+      input_padding: torch.LongTensor,
+      freq: torch.Tensor,
+  ) -> torch.Tensor:
+    num_outputs = len(self.config.quantiles) + 1
+    model_input, patched_padding, stats, _ = self._preprocess_input(
+        input_ts=input_ts,
+        input_padding=input_padding,
+    )
+    f_emb = self.freq_emb(freq)  # B x 1 x D
+    model_input += f_emb
+    model_output = self.stacked_transformer(model_input, patched_padding)
+
+    output_ts = self._postprocess_output(model_output, num_outputs, stats)
+    return output_ts
+
+  def decode(
+      self,
+      input_ts: torch.Tensor,
+      paddings: torch.Tensor,
+      freq: torch.LongTensor,
+      horizon_len: int,
+      output_patch_len: int | None = None,
+      max_len: int | None = None,
+      return_forecast_on_context: bool = False,
+  ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Auto-regressive decoding without caching.
+
+    Args:
+      input_ts: input time-series and paddings. Time-series shape B x C.
+      paddings: padding shape B x (C + H) where H is the prediction length.
+      freq: frequency shape B x 1
+      horizon_len: prediction length.
+      output_patch_len: output length to be fetched from one step of
+        auto-regressive decoding.
+      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:
+      Tuple of two forecasting results:
+      - Point (mean) output predictions as a tensor with shape B x H'.
+      - Full predictions (mean and quantiles) as a tensor with shape
+        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 = input_ts
+    context_len = final_out.shape[1]
+    full_outputs = []
+    if max_len is None:
+      max_len = context_len
+    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}")
+    if output_patch_len is None:
+      output_patch_len = self.config.horizon_len
+    num_decode_patches = (horizon_len + output_patch_len -
+                          1) // output_patch_len
+    for step_index in range(num_decode_patches):
+      current_padding = paddings[:, 0:final_out.shape[1]]
+      input_ts = final_out[:, -max_len:]
+      input_padding = current_padding[:, -max_len:]
+      fprop_outputs = self(input_ts, input_padding, freq)
+      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[:, 0:-1, 0:self.config.patch_len, :]
+        new_full_ts = new_full_ts.reshape(new_full_ts.size(0), -1,
+                                          new_full_ts.size(3))
+
+        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(new_full_ts)
+      final_out = torch.concatenate([final_out, new_ts], axis=-1)
+
+    if return_forecast_on_context:
+      # `full_outputs` indexing starts at after the first input patch.
+      full_outputs = torch.concatenate(
+          full_outputs,
+          axis=1)[:, :(context_len - self.config.patch_len + horizon_len), :]
+    else:
+      # `full_outputs` indexing starts at the forecast horizon.
+      full_outputs = torch.concatenate(full_outputs, axis=1)[:,
+                                                             0:horizon_len, :]
+
+    return (full_outputs[:, :, 0], full_outputs)

v1/src/timesfm/time_features.py

@@ -0,0 +1,215 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+
+"""Directory to extract time covariates.
+
+Extract time covariates from datetime.
+"""
+
+import numpy as np
+import pandas as pd
+from pandas.tseries.holiday import EasterMonday
+from pandas.tseries.holiday import GoodFriday
+from pandas.tseries.holiday import Holiday
+from pandas.tseries.holiday import SU
+from pandas.tseries.holiday import TH
+from pandas.tseries.holiday import USColumbusDay
+from pandas.tseries.holiday import USLaborDay
+from pandas.tseries.holiday import USMartinLutherKingJr
+from pandas.tseries.holiday import USMemorialDay
+from pandas.tseries.holiday import USPresidentsDay
+from pandas.tseries.holiday import USThanksgivingDay
+from pandas.tseries.offsets import DateOffset
+from pandas.tseries.offsets import Day
+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
+
+
+def _distance_to_holiday(holiday):
+  """Return distance to given holiday."""
+
+  def _distance_to_day(index):
+    holiday_date = holiday.dates(
+        index - pd.Timedelta(days=MAX_WINDOW),
+        index + pd.Timedelta(days=MAX_WINDOW),
+    )
+    assert (
+        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
+
+  return _distance_to_day
+
+
+EasterSunday = Holiday(
+    "Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
+)
+NewYearsDay = Holiday("New Years Day", month=1, day=1)
+SuperBowl = Holiday(
+    "Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1))
+)
+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)
+NewYearsEve = Holiday("New Years Eve", month=12, day=31)
+BlackFriday = Holiday(
+    "Black Friday",
+    month=11,
+    day=1,
+    offset=[pd.DateOffset(weekday=TH(4)), Day(1)],
+)
+CyberMonday = Holiday(
+    "Cyber Monday",
+    month=11,
+    day=1,
+    offset=[pd.DateOffset(weekday=TH(4)), Day(4)],
+)
+
+HOLIDAYS = [
+    EasterMonday,
+    GoodFriday,
+    USColumbusDay,
+    USLaborDay,
+    USMartinLutherKingJr,
+    USMemorialDay,
+    USPresidentsDay,
+    USThanksgivingDay,
+    EasterSunday,
+    NewYearsDay,
+    SuperBowl,
+    MothersDay,
+    IndependenceDay,
+    ChristmasEve,
+    ChristmasDay,
+    NewYearsEve,
+    BlackFriday,
+    CyberMonday,
+]
+
+
+class TimeCovariates(object):
+  """Extract all time covariates except for holidays."""
+
+  def __init__(
+      self,
+      datetimes,
+      normalized=True,
+      holiday=False,
+  ):
+    """Init function.
+
+    Args:
+      datetimes: pandas DatetimeIndex (lowest granularity supported is min)
+      normalized: whether to normalize features or not
+      holiday: fetch holiday features or not
+
+    Returns:
+      None
+    """
+    self.normalized = normalized
+    self.dti = datetimes
+    self.holiday = holiday
+
+  def _minute_of_hour(self):
+    minutes = np.array(self.dti.minute, dtype=np.float32)
+    if self.normalized:
+      minutes = minutes / 59.0 - 0.5
+    return minutes
+
+  def _hour_of_day(self):
+    hours = np.array(self.dti.hour, dtype=np.float32)
+    if self.normalized:
+      hours = hours / 23.0 - 0.5
+    return hours
+
+  def _day_of_week(self):
+    day_week = np.array(self.dti.dayofweek, dtype=np.float32)
+    if self.normalized:
+      day_week = day_week / 6.0 - 0.5
+    return day_week
+
+  def _day_of_month(self):
+    day_month = np.array(self.dti.day, dtype=np.float32)
+    if self.normalized:
+      day_month = day_month / 30.0 - 0.5
+    return day_month
+
+  def _day_of_year(self):
+    day_year = np.array(self.dti.dayofyear, dtype=np.float32)
+    if self.normalized:
+      day_year = day_year / 364.0 - 0.5
+    return day_year
+
+  def _month_of_year(self):
+    month_year = np.array(self.dti.month, dtype=np.float32)
+    if self.normalized:
+      month_year = month_year / 11.0 - 0.5
+    return month_year
+
+  def _week_of_year(self):
+    week_year = np.array(self.dti.strftime("%U").astype(int), dtype=np.float32)
+    if self.normalized:
+      week_year = week_year / 51.0 - 0.5
+    return week_year
+
+  def _get_holidays(self):
+    dti_series = self.dti.to_series()
+    hol_variates = np.vstack([
+        dti_series.apply(_distance_to_holiday(h)).values for h in tqdm(HOLIDAYS)
+    ])
+    # hol_variates is (num_holiday, num_time_steps), the normalization should be
+    # performed in the num_time_steps dimension.
+    return StandardScaler().fit_transform(hol_variates.T).T
+
+  def get_covariates(self):
+    """Get all time covariates."""
+    moh = self._minute_of_hour().reshape(1, -1)
+    hod = self._hour_of_day().reshape(1, -1)
+    dom = self._day_of_month().reshape(1, -1)
+    dow = self._day_of_week().reshape(1, -1)
+    doy = self._day_of_year().reshape(1, -1)
+    moy = self._month_of_year().reshape(1, -1)
+    woy = self._week_of_year().reshape(1, -1)
+
+    all_covs = [
+        moh,
+        hod,
+        dom,
+        dow,
+        doy,
+        moy,
+        woy,
+    ]
+    columns = ["moh", "hod", "dom", "dow", "doy", "moy", "woy"]
+    if self.holiday:
+      hol_covs = self._get_holidays()
+      all_covs.append(hol_covs)
+      columns += [f"hol_{i}" for i in range(len(HOLIDAYS))]
+
+    return pd.DataFrame(
+        data=np.vstack(all_covs).transpose(),
+        columns=columns,
+        index=self.dti,
+    )

v1/src/timesfm/timesfm_base.py

@@ -0,0 +1,736 @@
+# 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.
+"""Base class for TimesFM inference. This will be common to PAX and Pytorch."""
+
+import collections
+import dataclasses
+import logging
+import multiprocessing
+from typing import Any, Literal, Sequence, TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+from utilsforecast.processing import make_future_dataframe
+
+if TYPE_CHECKING:
+    from . import xreg_lib
+    Category = xreg_lib.Category
+    XRegMode = xreg_lib.XRegMode
+else:
+    Category = int | str
+    XRegMode = str
+
+_TOL = 1e-6
+DEFAULT_QUANTILES = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
+
+
+def process_group(key, group, value_name, forecast_context_len):
+  group = group.tail(forecast_context_len)
+  return np.array(group[value_name], dtype=np.float32), key
+
+
+def moving_average(arr, window_size):
+  """Calculates the moving average using NumPy's convolution function."""
+  # Pad with zeros to handle initial window positions
+  arr_padded = np.pad(arr, (window_size - 1, 0), "constant")
+  smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / 
+                  window_size)
+  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("MS"):
+    return 1
+  elif freq.endswith(("H", "T", "MIN", "D", "B", "U", "S")):
+    return 0
+  elif (
+    freq.endswith(("W", "M"))
+    or freq.startswith("W-")
+    or (freq.startswith("M") and len(freq) == 2)
+  ):
+    return 1
+  elif (
+    freq.endswith(("Y", "Q", "A"))
+    or freq.startswith("Y-")
+    or freq.startswith("Q-")
+    or freq.startswith("A-")
+  ):
+    return 2
+  else:
+    raise ValueError(f"Invalid frequency: {freq}")
+
+
+def strip_leading_nans(arr):
+  """
+  Removes contiguous NaN values from the beginning of a NumPy array.
+
+  Args:
+    arr: The input NumPy array.
+
+  Returns:
+    A new NumPy array with leading NaN values removed.
+    If the array is all NaNs or empty, returns an empty array.
+  """
+
+  isnan = np.isnan(arr)
+  first_valid_index = np.argmax(~isnan)
+  return arr[first_valid_index:]
+
+
+def linear_interpolation(arr):
+  """
+    Performs linear interpolation to fill NaN values in a 1D numpy array.
+
+    Args:
+        arr: The 1D numpy array containing NaN values.
+
+    Returns:
+        A new numpy array with NaN values filled using linear interpolation, 
+        or the original array if no NaNs are present. 
+        Returns None if the input is not a 1D array.
+        Returns the original array if there are no NaN values.
+    """
+
+  nans = np.isnan(arr)
+  if not np.any(nans):  # Check if there are any NaNs
+    return arr
+
+  def x(z):
+    return z.nonzero()[0]
+
+  nans_indices = x(nans)
+  non_nans_indices = x(~nans)
+  non_nans_values = arr[~nans]
+
+  try:
+    arr[nans] = np.interp(nans_indices, non_nans_indices, non_nans_values)
+  except ValueError:
+    if len(non_nans_values) > 0:
+      mu = np.nanmean(arr)
+    else:
+      mu = 0.0
+    arr = np.where(np.isfinite(arr), arr, mu)
+  return arr
+
+
+# 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)]
+
+
+@dataclasses.dataclass(kw_only=True)
+class TimesFmHparams:
+  """Hparams used to initialize a TimesFM model for inference.
+
+  These are the sufficient subset of hparams to configure TimesFM inference
+  agnostic to the checkpoint version, and are not necessarily the same as the
+  hparams used to train the checkpoint.
+
+  Attributes:
+    context_len: Largest context length the model allows for each decode call.
+      This technically can be any large, but practically should set to the
+      context length the checkpoint was trained with.
+    horizon_len: Forecast horizon.
+    input_patch_len: Input patch len.
+    output_patch_len: Output patch len. How many timepoints is taken from a
+      single step of autoregressive decoding. Can be set as the training horizon
+      of the checkpoint.
+    num_layers: Number of transformer layers in the model.
+    model_dims: Model dimension.
+    per_core_batch_size: Batch size on each core for data parallelism.
+    backend: One of "cpu", "gpu" or "tpu".
+    quantiles: Which quantiles are output by the model.
+  """
+
+  context_len: int = 512
+  horizon_len: int = 128
+  input_patch_len: int = 32
+  output_patch_len: int = 128
+  num_layers: int = 20
+  num_heads: int = 16
+  model_dims: int = 1280
+  per_core_batch_size: int = 32
+  backend: Literal["cpu", "gpu", "tpu"] = "cpu"
+  quantiles: Sequence[float] | None = DEFAULT_QUANTILES
+  use_positional_embedding: bool = True
+  # Hparams beyond the model.
+  point_forecast_mode: Literal["mean", "median"] = "median"
+
+
+@dataclasses.dataclass(kw_only=True)
+class TimesFmCheckpoint:
+  """Checkpoint used to initialize a TimesFM model for inference.
+
+  Attributes:
+    version: Version of the checkpoint, e.g. "jax", "torch", "tensorflow", etc.
+      The factory will create the corresponding TimesFm inference class based on
+      this version.
+    path: Path to the checkpoint.
+    type: If provided, type of the checkpoint used by the specific checkpoint
+      loader per version.
+    step: If provided, step of the checkpoint.
+  """
+
+  version: str = "jax"
+  path: str | None = None
+  huggingface_repo_id: str | None = None
+  type: Any = None
+  step: int | None = None
+  local_dir: str | None = None
+
+
+class TimesFmBase:
+  """Base TimesFM forecast API for inference.
+
+  This class is the scaffolding for calling TimesFM forecast. To properly use:
+    1. Create an instance with the correct hyperparameters of a TimesFM model.
+    2. Call `load_from_checkpoint` to load a compatible checkpoint.
+    3. Call `forecast` for inference.
+  """
+
+  def _logging(self, s):
+    print(s)
+
+  def __post_init__(self) -> None:
+    """Additional initialization for subclasses before checkpoint loading."""
+    pass
+
+  def __init__(self, hparams: TimesFmHparams,
+               checkpoint: TimesFmCheckpoint) -> None:
+    """Initializes the TimesFM forecast API.
+
+    Args:
+      hparams: Hyperparameters of the model.
+      checkpoint: Checkpoint to load. Notice `checkpoint.version` will decide
+        which TimesFM version to use.
+    """
+    self.hparams = hparams
+
+    # Expand hparams for conciseness within the model code.
+    self.context_len = hparams.context_len
+    self.horizon_len = hparams.horizon_len
+    self.input_patch_len = hparams.input_patch_len
+    self.output_patch_len = hparams.output_patch_len
+    self.num_layers = hparams.num_layers
+    self.model_dims = hparams.model_dims
+    self.backend = hparams.backend
+    self.quantiles = hparams.quantiles
+    self.num_heads = hparams.num_heads
+    self.use_pos_emb = hparams.use_positional_embedding
+
+    # Rewrite these values in __post_init__ for SPMD.
+    self.num_cores = 1
+    self.per_core_batch_size = hparams.per_core_batch_size
+    self.global_batch_size = hparams.per_core_batch_size
+
+    self._horizon_start = self.context_len - self.input_patch_len
+    self.__post_init__()
+    self.load_from_checkpoint(checkpoint)
+
+  def load_from_checkpoint(self, checkpoint: TimesFmCheckpoint) -> None:
+    """Loads a checkpoint and compiles the decoder."""
+    raise NotImplementedError("`load_from_checkpoint` is not implemented.")
+
+  def _preprocess(
+      self, inputs: Sequence[np.ndarray],
+      freq: Sequence[int]) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
+    """Formats and pads raw inputs to feed into the model.
+
+    This function both pads each time series to match the context length, and
+    pads the inputs to meet the SPMD shape requirement.
+
+    Args:
+      inputs: A list of 1d JTensors. Each JTensor is the context time series of
+        a single forecast task.
+      freq: list of frequencies
+
+    Returns:
+    A tuple of:
+    - the padded input time series to meet the model required context.
+    - the padding indicator.
+    - the frequency of each input time series.
+    - the number of padded examples for SPMD so that each core has the same
+        number (a multiple of `batch_size`) of examples.
+    """
+
+    input_ts, input_padding, inp_freq = [], [], []
+
+    pmap_pad = ((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],
+                            axis=0)
+        padding = np.concatenate(
+            [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):]
+
+      input_ts.append(ts)
+      input_padding.append(padding)
+      inp_freq.append(freq[i])
+
+    # Padding the remainder batch.
+    for _ in range(pmap_pad):
+      input_ts.append(input_ts[-1])
+      input_padding.append(input_padding[-1])
+      inp_freq.append(inp_freq[-1])
+
+    return (
+        np.stack(input_ts, axis=0),
+        np.stack(input_padding, axis=0),
+        np.array(inp_freq).astype(np.int32).reshape(-1, 1),
+        pmap_pad,
+    )
+
+  def _forecast(
+      self,
+      inputs: Sequence[Any],
+      freq: Sequence[int] | None = None,
+      window_size: int | None = None,
+      forecast_context_len: int | None = None,
+      return_forecast_on_context: bool = False,
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Forecasts on a list of time series.
+
+    Args:
+      inputs: list of time series forecast contexts. Each context time series
+        should be in a format convertible to JTensor by `jnp.array`.
+      freq: frequency of each context time series. 0 for high frequency
+        (default), 1 for medium, and 2 for low. Notice this is different from
+        the `freq` required by `forecast_on_df`.
+      window_size: window size of trend + residual decomposition. If None then
+        we do not do decomposition.
+      forecast_context_len: optional max context length.
+      return_forecast_on_context: True to return the forecast on the context
+        when available, i.e. after the first input patch.
+
+    Returns:
+    A tuple for np.array:
+    - the mean forecast of size (# inputs, # forecast horizon),
+    - the full forecast (mean + quantiles) of size
+        (# inputs,  # forecast horizon, 1 + # quantiles).
+
+    Raises:
+    ValueError: If the checkpoint is not properly loaded.
+    """
+    raise NotImplementedError("`_forecast` is not implemented.")
+
+  def forecast(
+      self,
+      inputs: Sequence[Any],
+      freq: Sequence[int] | None = None,
+      window_size: int | None = None,
+      forecast_context_len: int | None = None,
+      return_forecast_on_context: bool = False,
+      normalize: bool = False,
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Forecasts on a list of time series.
+
+    Args:
+      inputs: list of time series forecast contexts. Each context time series
+        should be in a format convertible to JTensor by `jnp.array`.
+      freq: frequency of each context time series. 0 for high frequency
+        (default), 1 for medium, and 2 for low. Notice this is different from
+        the `freq` required by `forecast_on_df`.
+      window_size: window size of trend + residual decomposition. If None then
+        we do not do decomposition.
+      forecast_context_len: optional max context length.
+      return_forecast_on_context: True to return the forecast on the context
+        when available, i.e. after the first input patch.
+      normalize: If True, then we normalize the inputs before forecasting and
+        the outputs are then renormalized to the original scale.
+
+    Returns:
+    A tuple for np.array:
+    - the mean forecast of size (# inputs, # forecast horizon),
+    - the full forecast (mean + quantiles) of size
+        (# inputs,  # forecast horizon, 1 + # quantiles).
+
+    Raises:
+    ValueError: If the checkpoint is not properly loaded.
+    """
+    stats = None
+
+    tmp_inputs = []
+    for each_input in inputs:
+      arr = np.array(each_input)
+      if not np.isfinite(arr).all():
+        arr = np.where(np.isfinite(arr), arr, np.nan)
+        arr = strip_leading_nans(arr)
+        arr = linear_interpolation(arr)
+      tmp_inputs.append(arr)
+
+    inputs = tmp_inputs
+    if normalize:
+      inputs, stats = _normalize(inputs)
+    mean_forecast, quantile_forecast = self._forecast(
+        inputs,
+        freq,
+        window_size,
+        forecast_context_len,
+        return_forecast_on_context,
+    )
+    if stats is not None:
+      stats = np.array(stats)
+      mu = stats[:, 0]
+      sigma = stats[:, 1]
+      mean_forecast = mean_forecast * sigma[:, None] + mu[:, None]
+      quantile_forecast = (quantile_forecast * sigma[:, None, None] +
+                           mu[:, None, None])
+    if self.hparams.point_forecast_mode == "mean":
+      return mean_forecast, quantile_forecast
+    elif self.hparams.point_forecast_mode == "median":
+      if self._median_index == -1:
+        for i, quantile in enumerate(self.quantiles):
+          if quantile == 0.5:
+            self._median_index = i
+            break
+        if self._median_index == -1:
+          raise ValueError("Median (0.5) is not found in the model quantiles:"
+                           f" {self.quantiles}. Please check the hparams.")
+      return (
+          quantile_forecast[:, :, 1 + self._median_index],
+          quantile_forecast,
+      )
+    else:
+      raise ValueError(
+          "Unsupported point forecast mode:"
+          f" {self.hparams.point_forecast_mode}. Use 'mean' or 'median'.")
+
+  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.
+    """
+
+    from . import xreg_lib
+
+    # Verify and bookkeep covariates.
+    if not (dynamic_numerical_covariates or dynamic_categorical_covariates or
+            static_numerical_covariates or static_categorical_covariates):
+      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,
+      freq: str,
+      forecast_context_len: int = 0,
+      value_name: str = "values",
+      model_name: str = "timesfm",
+      window_size: int | None = None,
+      num_jobs: int = 1,
+      normalize: bool = False,
+      verbose: bool = True,
+  ) -> pd.DataFrame:
+    """Forecasts on a list of time series.
+
+    Args:
+      inputs: A pd.DataFrame of all time series. The dataframe should have a
+        `unique_id` column for identifying the time series, a `ds` column for
+        timestamps and a value column for the time series values.
+      freq: string valued `freq` of data. Notice this is different from the
+        `freq` required by `forecast`. See `freq_map` for allowed values.
+      forecast_context_len: If provided none zero, we take the last
+        `forecast_context_len` time-points from each series as the forecast
+        context instead of the `context_len` set by the model.
+      value_name: The name of the value column.
+      model_name: name of the model to be written into future df.
+      window_size: window size of trend + residual decomposition. If None then
+        we do not do decomposition.
+      num_jobs: number of parallel processes to use for dataframe processing.
+      normalize: normalize context before forecasting or not.
+      verbose: output model states in terminal.
+
+    Returns:
+      Future forecasts dataframe.
+    """
+    if not ("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.")
+    if not forecast_context_len:
+      forecast_context_len = self.context_len
+    logging.info("Preprocessing dataframe.")
+    df_sorted = inputs.sort_values(by=["unique_id", "ds"])
+    new_inputs = []
+    uids = []
+    if num_jobs == 1:
+      if verbose:
+        print("Processing dataframe with single process.")
+      for key, group in df_sorted.groupby("unique_id"):
+        inp, uid = process_group(
+            key,
+            group,
+            value_name,
+            forecast_context_len,
+        )
+        new_inputs.append(inp)
+        uids.append(uid)
+    else:
+      if num_jobs == -1:
+        num_jobs = multiprocessing.cpu_count()
+      if verbose:
+        print("Processing dataframe with multiple processes.")
+      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")],
+        )
+      new_inputs, uids = zip(*results)
+    if verbose:
+      print("Finished preprocessing dataframe.")
+    freq_inps = [freq_map(freq)] * len(new_inputs)
+    _, full_forecast = self.forecast(new_inputs,
+                                     freq=freq_inps,
+                                     normalize=normalize,
+                                     window_size=window_size)
+    if verbose:
+      print("Finished forecasting.")
+    fcst_df = make_future_dataframe(
+        uids=uids,
+        last_times=df_sorted.groupby("unique_id")["ds"].tail(1),
+        h=self.horizon_len,
+        freq=freq,
+    )
+    fcst_df[model_name] = full_forecast[:, 0:self.horizon_len, 0].reshape(-1, 1)
+
+    for i, q in enumerate(self.quantiles):
+      q_col = f"{model_name}-q-{q}"
+      fcst_df[q_col] = full_forecast[:, 0:self.horizon_len,
+                                     1 + i].reshape(-1, 1)
+      if q == 0.5:
+        fcst_df[model_name] = fcst_df[q_col]
+    logging.info("Finished creating output dataframe.")
+    return fcst_df

v1/src/timesfm/timesfm_jax.py

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

v1/src/timesfm/timesfm_torch.py

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

v1/src/timesfm/xreg_lib.py

@@ -0,0 +1,486 @@
+# 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_output=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