Adding v2.0 support

src/timesfm/patched_decoder.py

@@ -333,7 +333,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     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])
@@ -401,7 +401,7 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
       inputs: NestedMap,
       horizon_len: int,
       output_patch_len: Optional[int] = None,
-      max_len: int = 512,
+      max_len: int | None = None,
       return_forecast_on_context: bool = False,
   ) -> tuple[JTensor, JTensor]:
     """Auto-regressive decoding without caching.
@@ -427,6 +427,8 @@ class PatchedTimeSeriesDecoder(base_layer.BaseLayer):
     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:

src/timesfm/pytorch_patched_decoder.py

@@ -730,7 +730,7 @@ class PatchedTimeSeriesDecoder(nn.Module):
       freq: torch.LongTensor,
       horizon_len: int,
       output_patch_len: int | None = None,
-      max_len: int = 512,
+      max_len: int | None = None,
       return_forecast_on_context: bool = False,
   ) -> tuple[torch.Tensor, torch.Tensor]:
     """Auto-regressive decoding without caching.
@@ -757,6 +757,8 @@ class PatchedTimeSeriesDecoder(nn.Module):
     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:"
@@ -773,9 +775,9 @@ class PatchedTimeSeriesDecoder(nn.Module):
       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.config.patch_len, :]
-        new_full_ts = fprop_outputs.view(new_full_ts.size(0), -1,
-                                         new_full_ts.size(3))
+        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)
 

src/timesfm/timesfm_base.py

@@ -61,6 +61,7 @@ def freq_map(freq: str):
   else:
     raise ValueError(f"Invalid frequency: {freq}")
 
+
 def strip_leading_nans(arr):
   """
   Removes contiguous NaN values from the beginning of a NumPy array.
@@ -77,6 +78,7 @@ def strip_leading_nans(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.
@@ -95,7 +97,9 @@ def linear_interpolation(arr):
   if not np.any(nans):  # Check if there are any NaNs
     return arr
 
-  x = lambda z: z.nonzero()[0]
+  def x(z):
+    return z.nonzero()[0]
+
   nans_indices = x(nans)
   non_nans_indices = x(~nans)
   non_nans_values = arr[~nans]
@@ -363,7 +367,7 @@ class TimesFmBase:
     ValueError: If the checkpoint is not properly loaded.
     """
     stats = None
-    
+
     tmp_inputs = []
     for each_input in inputs:
       arr = np.array(each_input)
@@ -372,7 +376,7 @@ class TimesFmBase:
         arr = strip_leading_nans(arr)
         arr = linear_interpolation(arr)
       tmp_inputs.append(arr)
-  
+
     inputs = tmp_inputs
     if normalize:
       inputs, stats = _normalize(inputs)

src/timesfm/timesfm_torch.py

@@ -58,8 +58,8 @@ class TimesFmTorch(timesfm_base.TimesFmBase):
     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")
+          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)
@@ -79,36 +79,33 @@ class TimesFmTorch(timesfm_base.TimesFmBase):
   ) -> 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.
+    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).
+    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.")
 
-        Raises:
-        ValueError: If the checkpoint is not properly loaded.
-        """
-    if 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]
+      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 = []
@@ -121,36 +118,39 @@ class TimesFmTorch(timesfm_base.TimesFmBase):
       freq = [0] * len(inputs)
 
     input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
+
     with torch.no_grad():
       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 = torch.from_numpy(
-            np.array(input_ts[i * self.global_batch_size:(i + 1) *
-                              self.global_batch_size],
-                     dtype=np.float32)).to(self._device)
-        input_padding_in = torch.from_numpy(
-            np.array(input_padding[i * self.global_batch_size:(i + 1) *
-                                   self.global_batch_size],
-                     dtype=np.float32)).to(self._device)
-        inp_freq_in = torch.from_numpy(
-            np.array(inp_freq[
-                i * self.global_batch_size:(i + 1) * self.global_batch_size,
-                :,
-            ],
-                     dtype=np.int32)).long().to(self._device)
+        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=input_ts_in,
-            paddings=input_padding_in,
-            freq=inp_freq_in,
+            input_ts=t_input_ts,
+            paddings=t_input_padding,
+            freq=t_inp_freq,
             horizon_len=self.horizon_len,
-            return_forecast_on_context=return_forecast_on_context,
+            output_patch_len=self.output_patch_len,
+            # Returns forecasts on context for parity with the Jax version.
+            return_forecast_on_context=True,
         )
-        mean_output = mean_output.detach().cpu().numpy()
-        full_output = full_output.detach().cpu().numpy()
-        mean_output = np.array(mean_output)
-        full_output = np.array(full_output)
+        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)
 
@@ -164,4 +164,5 @@ class TimesFmTorch(timesfm_base.TimesFmBase):
     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