Add troubleshooting section to readme file

src/adapter/__init__.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """adapter init file."""
 
 from .dora_layers import DoraAttentionProjection, DoraCombinedQKVProjection, DoraLinear

src/adapter/dora_layers.py

@@ -21,182 +21,181 @@ 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 __init__(self, module):
+    self.module = module
 
-    def _dorafy_var(self, w):
-        lora_a = super().__getattr__("lora_a")
-        lora_b = super().__getattr__("lora_b")
-        dora_m = super().__getattr__("dora_m")
+  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
 
-        lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
-        lora_delta = jnp.reshape(lora_delta, w.shape)
+  def _dorafy_var(self, w):
+    lora_a = super().__getattr__("lora_a")
+    lora_b = super().__getattr__("lora_b")
+    dora_m = super().__getattr__("dora_m")
 
-        w_prime = w + lora_delta
+    lora_delta = self.module.einsum("...dr,...nr->...dn", lora_a, lora_b)
+    lora_delta = jnp.reshape(lora_delta, w.shape)
 
-        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
+    w_prime = w + lora_delta
 
-    def __getattr__(self, k):
-        var = super().__getattr__(k)
-        if not self._dora_initialized():
-            return var
+    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
 
-        if k == "w":
-            return self._dorafy_var(var)
+  def __getattr__(self, k):
+    var = super().__getattr__(k)
+    if not self._dora_initialized():
+      return var
 
-        return var
+    if k == "w":
+      return self._dorafy_var(var)
 
-    def __getitem__(self, k):
-        var = super().__getattr__(k)
-        if not self._dora_initialized():
-            return var
+    return var
 
-        if k == "w":
-            return self._dorafy_var(var)
+  def __getitem__(self, k):
+    var = super().__getattr__(k)
+    if not self._dora_initialized():
+      return var
 
-        return var
+    if k == "w":
+      return self._dorafy_var(var)
+
+    return var
 
 
 class DoraThetaDescriptor:
-    """Dot syntax accession descriptor."""
+  """Dot syntax accession descriptor."""
 
-    def __get__(self, obj, objtype=None):
-        return DoraTheta(obj)
+  def __get__(self, obj, objtype=None):
+    return DoraTheta(obj)
 
 
 class DoraLinear(linears.Linear):
-    rank: int = 0
-    lora_init: WeightInit | None = None
-    theta = DoraThetaDescriptor()
+  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
+  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],
-            ),
-        )
+    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()
+  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
+  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],
-            ),
-        )
+    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()
+  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
+  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],
-            ),
-        )
+    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],
+        ),
+    )

src/adapter/lora_layers.py

@@ -21,146 +21,145 @@ 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 __init__(self, module):
+    self.module = module
 
-    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 _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 __getattr__(self, k):
-        var = super().__getattr__(k)
-        if not self._lora_initialized():
-            return var
+  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
 
-        if k == "w":
-            return self._lorafy_var(var)
+  def __getattr__(self, k):
+    var = super().__getattr__(k)
+    if not self._lora_initialized():
+      return var
 
-        return var
+    if k == "w":
+      return self._lorafy_var(var)
 
-    def __getitem__(self, k):
-        var = super().__getattr__(k)
-        if not self._lora_initialized():
-            return var
+    return var
 
-        if k == "w":
-            return self._lorafy_var(var)
+  def __getitem__(self, k):
+    var = super().__getattr__(k)
+    if not self._lora_initialized():
+      return var
 
-        return var
+    if k == "w":
+      return self._lorafy_var(var)
+
+    return var
 
 
 class LoraThetaDescriptor:
-    """Dot syntax accession descriptor."""
+  """Dot syntax accession descriptor."""
 
-    def __get__(self, obj, objtype=None):
-        return LoraTheta(obj)
+  def __get__(self, obj, objtype=None):
+    return LoraTheta(obj)
 
 
 class LoraLinear(linears.Linear):
-    rank: int = 0
-    lora_init: WeightInit | None = None
-    theta = LoraThetaDescriptor()
+  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
+  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],
-            ),
-        )
+    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()
+  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
+  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(
+        "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()
+  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
+  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(
+        "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],
+        ),
+    )

src/adapter/utils.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """
 This file provides functionality for loading and merging adapter weights
 in timesfm model, specifically for LoRA and DoRA.
@@ -40,10 +39,11 @@ from adapter.lora_layers import (
 from timesfm import TimesFm
 
 
-def get_adapter_params(
-    params: dict, lora_target_modules: str, num_layers: int, use_dora: bool = False
-) -> dict:
-    """
+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:
@@ -55,47 +55,44 @@ def get_adapter_params(
     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] = {}
+  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"]
+    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"]
+        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,
-                }
+        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 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"]
+    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"]
+      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,
-                }
+        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
+        if use_dora:
+          adapter_params[layer_key][component]["dora_m"] = attention[component][
+              "dora_m"]
+  return adapter_params
 
 
 def load_adapter_checkpoint(
@@ -105,7 +102,7 @@ def load_adapter_checkpoint(
     lora_target_modules: str,
     use_dora: bool,
 ) -> None:
-    """
+  """
     Loads an adapter checkpoint and merges it with the original model weights.
 
     Args:
@@ -118,83 +115,77 @@ def load_adapter_checkpoint(
     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,
-        )
-    )
+  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
-    )
+  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_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,
-    )
+  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,
-    )
+  # 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
-        )
+  # 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."
-    )
+  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()
+  # jit compile the model
+  model.jit_decode()
 
 
 def _merge_adapter_weights(
@@ -204,7 +195,7 @@ def _merge_adapter_weights(
     num_layers: int,
     use_dora: bool,
 ) -> None:
-    """
+  """
     Merges adapter weights with the original model weights.
 
     Args:
@@ -214,74 +205,73 @@ def _merge_adapter_weights(
         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}"
+  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"]
+    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"]
+        params = adapter_train_state.mdl_vars[layer_key][ff_layer_key]
+        lora_a = params["lora_a"]
+        lora_b = params["lora_b"]
 
-                w = linear["w"]
+        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
+        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"]
+        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
+        else:
+          linear["w"] = w_prime
 
-                del linear["lora_a"]
-                del linear["lora_b"]
+        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"]
+    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"]
+      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"]
+        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
+        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"]
+        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
+        else:
+          attention[component]["w"] = w_prime
 
-                del attention[component]["lora_a"]
-                del attention[component]["lora_b"]
+        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:
-    """
+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:
@@ -293,42 +283,39 @@ def _get_adapter_weight_params(
     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] = {}
+  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 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 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 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"
-                    ]
+        if use_dora:
+          adapter_params[layer][component]["dora_m"] = adapter_weight_params[
+              "dora_m"]
 
-    return adapter_params
+  return adapter_params
 
 
 def load_adapter_layer(
@@ -338,7 +325,7 @@ def load_adapter_layer(
     lora_target_modules: str,
     use_dora: bool = False,
 ) -> tuple[pax_fiddle.Config, pax_fiddle.Config]:
-    """
+  """
     Updates target modules with adapter layers.
 
     Args:
@@ -351,67 +338,55 @@ def load_adapter_layer(
     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
-        )
+  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
-        )
+  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)
+    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
-        )
+    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)
+    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
-        )
+    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,
-    )
+  # 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
+  return original_linear_tpl, original_attn_tpl, original_combined_qkv_tpl
 
 
 def _initialize_adapter_params(
@@ -422,7 +397,7 @@ def _initialize_adapter_params(
     use_dora: bool = False,
     seed: int = 1234,
 ) -> dict:
-    """
+  """
     Initializes and adds adapter parameters to target modules.
 
     Args:
@@ -436,52 +411,47 @@ def _initialize_adapter_params(
     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)
+  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))
+        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
+        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 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"
-            ]
+    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)
+      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))
+        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
+        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
+        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

src/finetuning/finetuning_example.py

@@ -43,11 +43,11 @@ flags.DEFINE_list(
 )
 
 flags.DEFINE_string(
-    "local_model_path",
-    None,
+    "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."""
 
@@ -148,7 +148,7 @@ def get_model(load_weights: bool = False):
       use_positional_embedding=False,
       context_len=192,
   )
-  
+
   if load_weights:
     if FLAGS.local_model_path:
       tfm_config = TimesFMConfig()
@@ -157,11 +157,12 @@ def get_model(load_weights: bool = False):
     else:
       repo_id = "google/timesfm-2.0-500m-pytorch"
       tfm = TimesFm(hparams=hparams,
-              checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
+                    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")
+      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)

src/timesfm/__init__.py

@@ -25,11 +25,13 @@ from timesfm.timesfm_base import (
 import sys
 
 try:
-    from timesfm.timesfm_jax import TimesFmJax as TimesFm
-    from timesfm import data_loader
+  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}.")
+  print(f"Loaded Jax TimesFM, likely because python version is {sys.version}.")
 except Exception as _:
-    from timesfm.timesfm_torch import TimesFmTorch as TimesFm
+  from timesfm.timesfm_torch import TimesFmTorch as TimesFm
 
-    print(f"Loaded PyTorch TimesFM, likely because python version is {sys.version}.")
+  print(
+      f"Loaded PyTorch TimesFM, likely because python version is {sys.version}."
+  )

src/timesfm/time_features.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Directory to extract time covariates.
 
 Extract time covariates from datetime.
@@ -36,7 +35,6 @@ from pandas.tseries.offsets import Easter
 from sklearn.preprocessing import StandardScaler
 from tqdm import tqdm
 
-
 # This is 183 to cover half a year (in both directions), also for leap years
 # + 17 as Eastern can be between March, 22 - April, 25
 MAX_WINDOW = 183 + 17
@@ -50,9 +48,8 @@ def _distance_to_holiday(holiday):
         index - pd.Timedelta(days=MAX_WINDOW),
         index + pd.Timedelta(days=MAX_WINDOW),
     )
-    assert (
-        len(holiday_date) != 0  # pylint: disable=g-explicit-length-test
-    ), f"No closest holiday for the date index {index} found."
+    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
@@ -60,16 +57,19 @@ def _distance_to_holiday(holiday):
   return _distance_to_day
 
 
-EasterSunday = Holiday(
-    "Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)]
-)
+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))
-)
+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)

src/timesfm/timesfm_base.py

@@ -25,12 +25,12 @@ 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
+  from . import xreg_lib
+  Category = xreg_lib.Category
+  XRegMode = xreg_lib.XRegMode
 else:
-    Category = int | str
-    XRegMode = str
+  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)
@@ -45,7 +45,7 @@ def moving_average(arr, window_size):
   """Calculates the moving average using NumPy's convolution function."""
   # Pad with zeros to handle initial window positions
   arr_padded = np.pad(arr, (window_size - 1, 0), "constant")
-  smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / 
+  smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") /
                   window_size)
   return [smoothed_arr, arr - smoothed_arr]
 
@@ -57,18 +57,11 @@ def freq_map(freq: str):
     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)
-  ):
+  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-")
-  ):
+  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}")