From 78fcba13ccceea3de81f3a47a124d60b8b8fc43a Mon Sep 17 00:00:00 2001
From: Rajat Sen <rsen91@gmail.com>
Date: Tue, 27 Aug 2024 18:49:28 +0000
Subject: [PATCH]  Revert "Merge pull request #130 from
 google-research/rajat_dev"

This reverts commit 9d4be94d14135cfcf140fa995c707f788a0552a5, reversing
changes made to dbe0fa528d58aa2946a88125a102a3417fe664b6.
---
 src/timesfm_torch/README.md                  |  10 -
 src/timesfm_torch/pytorch_patched_decoder.py | 811 -------------------
 2 files changed, 821 deletions(-)
 delete mode 100644 src/timesfm_torch/README.md
 delete mode 100644 src/timesfm_torch/pytorch_patched_decoder.py

diff --git a/src/timesfm_torch/README.md b/src/timesfm_torch/README.md
deleted file mode 100644
index 7bd28d7..0000000
--- a/src/timesfm_torch/README.md
+++ /dev/null
@@ -1,10 +0,0 @@
-# Official Pytorch implementation of TimesFM
-
-TimesFM  (Time Series Foundation Model) is a pretrained time-series foundation model developed by Google
-Research for time-series forecasting.
-
-* Paper: [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688), to appear in ICML 2024.
-* [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/)
-* [Hugging Face checkpoint repo](https://huggingface.co/google/timesfm-1.0-200m)
-
-## Please stay tuned for all the same functionalities as the pax version.
diff --git a/src/timesfm_torch/pytorch_patched_decoder.py b/src/timesfm_torch/pytorch_patched_decoder.py
deleted file mode 100644
index 9032692..0000000
--- a/src/timesfm_torch/pytorch_patched_decoder.py
+++ /dev/null
@@ -1,811 +0,0 @@
-# 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.25, 0.5, 0.75, 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. This should reproduce
-  http://google3/learning/multipod/pax/foundation_ts/patched_decoder.py?cl=664903156&l=320
-
-  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.
-  """
-  # TODO(senrajat): will be tested as a part of preprocess input directly.
-  # 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.
-
-  This should reproduce:
-  https://google3/learning/multipod/pax/foundation_ts/patched_decoder.py;l=125.
-
-  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.
-
-  It should mimic:
-  https://google3/third_party/py/praxis/layers/attentions.py;l=191
-
-  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[:, 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.
-
-  It should mimic:
-  https://source.corp.google.com/piper///depot/google3/third_party/py/praxis/layers/attentions.py;l=77
-
-  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.
-  It should mimic:
-  https://google3/third_party/py/praxis/layers/attentions.py;l=127
-
-  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 = 512,
-      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 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[:, :-1, :self.config.patch_len, :]
-        new_full_ts = fprop_outputs.view(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)