To comply with Google OSS style guide.

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -56,12 +56,10 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
 
     # Layers.
     self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
-        self.stacked_xf = nn.ModuleList(
+    self.stacked_xf = nn.ModuleList([
-            [
         transformer.Transformer(self.config.stacked_transformers.transformer)
         for _ in range(self.x)
-            ]
+    ])
-        )
     self.output_projection_point = dense.ResidualBlock(
         self.config.output_projection_point
     )
@@ -182,7 +180,9 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
           (batch_size, -1, self.o, self.q),
       )
       renormed_quantile_spread = torch.reshape(
-                revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
+          revin(
+              normed_quantile_spread, context_mu, context_sigma, reverse=True
+          ),
           (batch_size, -1, self.os, self.q),
       )[:, -1, ...]
 
@@ -266,11 +266,19 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
       path: str | None = None,
       hf_repo_id: str | None = "google/timesfm-2.5-200m-pytorch",
   ):
-        """Loads a PyTorch safetensors TimesFM model."""
+    """Loads a PyTorch safetensors TimesFM model.
+
+    Args:
+        path: Path to a local checkpoint. If not provided, will try to download
+          from the default Hugging Face repo.
+        hf_repo_id: Use another Hugging Face repo ID.
+    """
     if path:
       pass
     elif hf_repo_id:
-            logging.info("Downloading checkpoint from Hugging Face repo %s", hf_repo_id)
+      logging.info(
+          "Downloading checkpoint from Hugging Face repo %s", hf_repo_id
+      )
       path = os.path.join(
           huggingface_hub.snapshot_download(hf_repo_id), "model.safetensors"
       )
@@ -294,7 +302,8 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
           "When compiling, max context needs to be multiple of the patch size"
           " %d. Using max context = %d instead.",
           self.model.p,
-                new_context := math.ceil(fc.max_context / self.model.p) * self.model.p,
+          new_context := math.ceil(fc.max_context / self.model.p)
+          * self.model.p,
       )
       fc.max_context = new_context
     if fc.max_horizon % self.model.o != 0:
@@ -302,7 +311,8 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5):
           "When compiling, max horizon needs to be multiple of the output patch"
           " size %d. Using max horizon = %d instead.",
           self.model.o,
-                new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o,
+          new_horizon := math.ceil(fc.max_horizon / self.model.o)
+          * self.model.o,
       )
       fc.max_horizon = new_horizon
     if fc.max_context + fc.max_horizon > self.model.config.context_limit:

src/timesfm/torch/transformer.py

@@ -86,7 +86,8 @@ class RotaryPositionalEmbedding(nn.Module):
         / self.embedding_dims
     )
     timescale = (
-            self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
+        self.min_timescale
+        * (self.max_timescale / self.min_timescale) ** fraction
     ).to(inputs.device)
     if position is None:
       seq_length = inputs.shape[1]
@@ -208,16 +209,21 @@ class MultiHeadAttention(nn.Module):
           b, n_patches, dtype=torch.bool, device=inputs_q.device
       )
 
-        query = self.query(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+    query = self.query(inputs_q).view(
+        b, n_patches, self.num_heads, self.head_dim
+    )
     key = self.key(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
-        value = self.value(inputs_q).view(b, n_patches, self.num_heads, self.head_dim)
+    value = self.value(inputs_q).view(
+        b, n_patches, self.num_heads, self.head_dim
+    )
 
     if decode_cache is None:
       num_masked = torch.sum(patch_mask.to(torch.int32), dim=-1)
       next_index = torch.zeros_like(num_masked, dtype=torch.int32)
     else:
       num_masked = (
-                torch.sum(patch_mask.to(torch.int32), dim=-1) + decode_cache.num_masked
+          torch.sum(patch_mask.to(torch.int32), dim=-1)
+          + decode_cache.num_masked
       )
       next_index = decode_cache.next_index.clone()
 

src/timesfm/torch/util.py

@@ -49,7 +49,9 @@ def update_running_stats(
   inc_mu = inc_mu_numerator / inc_n_safe
   inc_mu = torch.where(inc_n == 0, 0.0, inc_mu)
 
-    inc_var_numerator = torch.sum(((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1)
+  inc_var_numerator = torch.sum(
+      ((x - inc_mu.unsqueeze(-1)) ** 2) * is_legit, dim=-1
+  )
   inc_var = inc_var_numerator / inc_n_safe
   inc_var = torch.where(inc_n == 0, 0.0, inc_var)
   inc_sigma = torch.sqrt(inc_var)