Merge pull request #323 from google-research/siriuz42-2.0-pr

flax

README.md

@@ -44,14 +44,29 @@ will be under construction over the next few weeks to
 
 ### Install
 
-TODO(siriuz42): Package timesfm==2.0.0 and upload to PyPI .
+**Step 1** (Prerequisite)
 
-Run
+Install your preferred `torch` / `jax` backend based on your OS and accelerators
+(CPU, GPU, TPU or Apple Silicon).
+
+-   [Install PyTorch](https://pytorch.org/get-started/locally/).
+-   [Install Jax](https://docs.jax.dev/en/latest/installation.html#installation)
+    for Flax.
+
+**Step 2** (Install `timesfm`)
+
+You can now install locally.
 
 ```shell
 git clone https://github.com/google-research/timesfm.git
 cd timesfm
-pip install -e .
+pip install -e ".[torch]"
+```
+
+```shell
+git clone https://github.com/google-research/timesfm.git
+cd timesfm
+pip install -e ".[flax]"
 ```
 
 ### Code Example

pyproject.toml

@@ -14,11 +14,22 @@ readme = "README.md"
 packages = [{include = "timesfm", from = "src"}]
 
 [tool.poetry.dependencies]
-python = ">=3.11"
+python = ">=3.11,<4.0"
 numpy = ">=1.26.4"
 huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
 safetensors = ">=0.5.3"
-torch = { version = ">=2.0.0", extras = ["cuda"] }
+torch = { version = ">=2.0.0", extras = ["cuda"], optional = true }
+flax = { version = ">=0.8.2", optional = true }
+optax = { version = ">=0.2.2", optional = true }
+einshape = { version = ">=0.8.0", optional = true }
+orbax-checkpoint = { version = ">=0.5.15", optional = true }
+jaxtyping = { version = ">=0.2.29", optional = true }
+jax = { version = ">=0.4.26", optional = true }
+
+
+[tool.poetry.extras]
+torch = ["torch"]
+flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "jax"]
 
 [tool.ruff]
 line-length = 88

src/timesfm/__init__.py

@@ -15,6 +15,15 @@
 """TimesFM API."""
 
 from .configs import ForecastConfig
-from .timesfm_2p5 import timesfm_2p5_torch
 
-TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch
+try:
+  from .timesfm_2p5 import timesfm_2p5_torch
+  TimesFM_2p5_200M_torch = timesfm_2p5_torch.TimesFM_2p5_200M_torch
+except ImportError:
+  pass
+
+try:
+  from .timesfm_2p5 import timesfm_2p5_flax
+  TimesFM_2p5_200M_flax = timesfm_2p5_flax.TimesFM_2p5_200M_flax
+except ImportError:
+  pass

src/timesfm/configs.py

@@ -18,7 +18,7 @@ import dataclasses
 from typing import Literal
 
 
-@dataclasses.dataclass(frozen=False)
+@dataclasses.dataclass(frozen=True)
 class ForecastConfig:
   """Options for forecasting.
 

src/timesfm/flax/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2025 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.

src/timesfm/flax/dense.py

@@ -0,0 +1,110 @@
+# Copyright 2025 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.
+
+"""Dense layers for TimesFM."""
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import jaxtyping
+
+from .. import configs
+
+Array = jaxtyping.Array
+Bool = jaxtyping.Bool
+Float = jaxtyping.Float
+Integer = jaxtyping.Integer
+Num = jaxtyping.Num
+
+ResidualBlockConfig = configs.ResidualBlockConfig
+RandomFourierFeaturesConfig = configs.RandomFourierFeaturesConfig
+
+
+class ResidualBlock(nnx.Module):
+  """Residual block with two linear layers and a linear residual connection."""
+
+  def __init__(self, config: ResidualBlockConfig, *, rngs=nnx.Rngs(42)):
+    self.config = config
+    self.hidden_layer = nnx.Linear(
+      in_features=config.input_dims,
+      out_features=config.hidden_dims,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+    self.output_layer = nnx.Linear(
+      in_features=config.hidden_dims,
+      out_features=config.output_dims,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+    self.residual_layer = nnx.Linear(
+      in_features=config.input_dims,
+      out_features=config.output_dims,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+    if config.activation == "relu":
+      self.activation = jax.nn.relu
+    elif config.activation == "swish":
+      self.activation = jax.nn.swish
+    elif config.activation == "none":
+      self.activation = lambda x: x
+    else:
+      raise ValueError(f"Activation: {config.activation} not supported.")
+
+  def __call__(self, x: Float[Array, "b ... i"]) -> Float[Array, "b ... o"]:
+    return self.output_layer(
+      self.activation(self.hidden_layer(x))
+    ) + self.residual_layer(x)
+
+
+class RandomFourierFeatures(nnx.Module):
+  """Random Fourier features layer."""
+
+  __data__ = ("phrase_shifts",)
+
+  def __init__(self, config: RandomFourierFeaturesConfig, *, rngs=nnx.Rngs(42)):
+    self.config = config
+
+    if config.output_dims % 4 != 0:
+      raise ValueError(
+        f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0."
+      )
+    num_projected_features = config.output_dims // 4
+
+    self.phase_shifts = nnx.Param(jnp.zeros(shape=(2, num_projected_features)))
+    self.projection_layer = nnx.Linear(
+      in_features=config.input_dims,
+      out_features=num_projected_features,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+    self.residual_layer = nnx.Linear(
+      in_features=config.input_dims,
+      out_features=config.output_dims,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+
+  def __call__(self, x: Float[Array, "b ... i"]) -> Float[Array, "b ... o"]:
+    projected = self.projection_layer(x)
+    cos_features = jnp.cos(projected)
+    sin_features = jnp.sin(projected)
+    sq_wave_1 = jnp.sign(jnp.sin(projected + self.phase_shifts[0, :]))
+    sq_wave_2 = jnp.sign(jnp.sin(projected + self.phase_shifts[1, :]))
+    fourier_features = jnp.concatenate(
+      [cos_features, sin_features, sq_wave_1, sq_wave_2], axis=-1
+    )
+    residual = self.residual_layer(x)
+    return fourier_features + residual

src/timesfm/flax/normalization.py

@@ -0,0 +1,71 @@
+# Copyright 2025 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.
+
+"""Normalization layers for TimesFM."""
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import jaxtyping
+
+Array = jaxtyping.Array
+Bool = jaxtyping.Bool
+Float = jaxtyping.Float
+Integer = jaxtyping.Integer
+Num = jaxtyping.Num
+
+
+class RMSNorm(nnx.Module):
+  """RMS normalization."""
+
+  __data__ = ("scale",)
+
+  def __init__(
+    self,
+    num_features: int,
+    *,
+    epsilon: float = 1e-6,
+    rngs=nnx.Rngs(42),
+  ):
+    del rngs
+    self.scale = nnx.Param(jnp.zeros(shape=(num_features,)))
+    self.num_features = num_features
+    self.epsilon = epsilon
+
+  def __call__(self, inputs: Float[Array, "b ... d"]) -> Float[Array, "b ... d"]:
+    var = jnp.mean(jnp.square(inputs), axis=-1, keepdims=True)
+    normed_inputs = inputs * jax.lax.rsqrt(var + self.epsilon)
+    normed_inputs *= self.scale
+    return normed_inputs
+
+
+class LayerNorm(nnx.Module):
+  """Layer normalization replica of  LayerNorm."""
+
+  __data__ = ("scale", "bias")
+
+  def __init__(self, num_features: int, *, epsilon: float = 1e-6, rngs=nnx.Rngs(42)):
+    del rngs
+    self.scale = nnx.Param(jnp.ones(shape=(num_features,)))
+    self.bias = nnx.Param(jnp.zeros(shape=(num_features,)))
+    self.num_features = num_features
+    self.epsilon = epsilon
+
+  def __call__(self, inputs: Float[Array, "b ... d"]) -> Float[Array, "b ... d"]:
+    mean = jnp.mean(inputs, axis=-1, keepdims=True)
+    var = jnp.mean(jnp.square(inputs - mean), axis=-1, keepdims=True)
+    normed_inputs = (inputs - mean) * jax.lax.rsqrt(var + self.epsilon)
+    normed_inputs *= self.scale
+    normed_inputs += self.bias
+    return normed_inputs

src/timesfm/flax/transformer.py

@@ -0,0 +1,356 @@
+# Copyright 2025 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.
+
+"""Transformer layers for TimesFM."""
+
+import functools
+from typing import Callable
+
+from flax import nnx
+from flax.nnx.nn import linear
+import jax
+from jax import lax
+import jax.numpy as jnp
+import jaxtyping
+
+from .. import configs
+from . import normalization, util
+
+Array = jaxtyping.Array
+Bool = jaxtyping.Bool
+Float = jaxtyping.Float
+Integer = jaxtyping.Integer
+Num = jaxtyping.Num
+LayerNorm = normalization.LayerNorm
+RMSNorm = normalization.RMSNorm
+LinearGeneral = linear.LinearGeneral
+TransformerConfig = configs.TransformerConfig
+DecodeCache = util.DecodeCache
+
+
+@functools.partial(
+  jax.jit,
+  static_argnames=("query_length", "kv_length"),
+)
+def make_attn_mask(
+  query_length: int,
+  num_all_masked_kv: Integer[Array, "b"],
+  query_index_offset: Integer[Array, "b"] | None = None,
+  kv_length: int = 0,
+) -> Bool[Array, "b 1 q n"]:
+  """Makes attention mask."""
+
+  if kv_length == 0:
+    kv_length = query_length
+
+  q_index = jnp.arange(query_length)[None, None, :, None]
+  if query_index_offset is not None:
+    q_index += query_index_offset[:, None, None, None]
+  kv_index = jnp.arange(kv_length)[None, None, None, :]
+  return jnp.logical_and(
+    q_index >= kv_index,
+    kv_index >= num_all_masked_kv[:, None, None, None],
+  )
+
+
+class RotaryPositionalEmbedding(nnx.Module):
+  """Rotary positional embedding."""
+
+  def __init__(
+    self,
+    embedding_dims: int,
+    min_timescale: int = 1,
+    max_timescale: int = 10000,
+  ):
+    self.embedding_dims = embedding_dims
+    self.min_timescale = min_timescale
+    self.max_timescale = max_timescale
+
+  def __call__(
+    self,
+    inputs: Float[Array, "b ... d"],
+    position: Array | None = None,
+  ):
+    """Generates a JTensor of sinusoids with different frequencies."""
+    if self.embedding_dims != inputs.shape[-1]:
+      raise ValueError(
+        "The embedding dims of the rotary position embedding"
+        "must match the hidden dimension of the inputs."
+      )
+    half_embedding_dim = self.embedding_dims // 2
+    fraction = 2 * jnp.arange(0, half_embedding_dim) / self.embedding_dims
+    timescale = (
+      self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
+    )
+    if position is None:
+      seq_length = inputs.shape[1]
+      position = jnp.arange(seq_length, dtype=jnp.float32)[None, :]
+    if len(inputs.shape) == 4:
+      position = position[..., None, None]
+      timescale = timescale[None, None, None, :]
+    elif len(inputs.shape) == 3:
+      position = position[..., None]
+      timescale = timescale[None, None, :]
+    else:
+      raise ValueError("Inputs must be of rank 3 or 4.")
+    sinusoid_inp = position / timescale
+    sin = jnp.sin(sinusoid_inp)
+    cos = jnp.cos(sinusoid_inp)
+    first_half, second_half = jnp.split(inputs, 2, axis=-1)
+    first_part = first_half * cos - second_half * sin
+    second_part = second_half * cos + first_half * sin
+    first_part = first_part.astype(None)
+    second_part = second_part.astype(None)
+    return jnp.concatenate([first_part, second_part], axis=-1)
+
+
+class PerDimScale(nnx.Module):
+  """Per-dimension scaling."""
+
+  __data__ = ("per_dim_scale",)
+
+  def __init__(self, num_dims: int, *, rngs=nnx.Rngs(42)):
+    del rngs
+    self.num_dims = num_dims
+    self.per_dim_scale = nnx.Param(jnp.zeros(shape=(num_dims,)))
+
+  def __call__(self, x: Float[Array, "b ... d"]) -> Float[Array, "b ... d"]:
+    return x * (
+      1.442695041 / jnp.sqrt(self.num_dims) * jax.nn.softplus(self.per_dim_scale)
+    )
+
+
+class MultiHeadAttention(nnx.Module):
+  """Multi-head attention."""
+
+  def __init__(
+    self,
+    num_heads: int,
+    in_features: int,
+    *,
+    use_per_dim_scale: bool = True,
+    use_rotary_position_embeddings: bool = True,
+    use_bias: bool = False,
+    deterministic: bool | None = None,
+    attention_fn: Callable[..., Array] = nnx.dot_product_attention,
+    qk_norm: str = "rms",
+    rngs=nnx.Rngs(42),
+  ):
+    self.num_heads = num_heads
+    self.in_features = in_features
+    self.qkv_features = in_features
+    self.out_features = in_features
+    self.in_kv_features = in_features
+    self.deterministic = deterministic
+    self.use_bias = use_bias
+    self.attention_fn = attention_fn
+    self.qk_norm = qk_norm
+
+    if self.qkv_features % self.num_heads != 0:
+      raise ValueError(
+        f"Memory dimension ({self.qkv_features}) must be divisible by "
+        f"'num_heads' heads ({self.num_heads})."
+      )
+    self.head_dim = self.qkv_features // self.num_heads
+
+    linear_general = functools.partial(
+      LinearGeneral,
+      out_features=(self.num_heads, self.head_dim),
+      use_bias=self.use_bias,
+    )
+    # project inputs_q to multi-headed q/k/v
+    # dimensions are then [batch..., length, n_heads, n_features_per_head]
+    self.query = linear_general(self.in_features, rngs=rngs)
+    self.key = linear_general(self.in_kv_features, rngs=rngs)
+    self.value = linear_general(self.in_kv_features, rngs=rngs)
+
+    if self.qk_norm == "rms":
+      self.query_ln = RMSNorm(self.head_dim)
+      self.key_ln = RMSNorm(self.head_dim)
+    else:
+      self.query_ln = None
+      self.key_ln = None
+
+    self.out = LinearGeneral(
+      in_features=(self.num_heads, self.head_dim),
+      out_features=self.out_features,
+      axis=(-2, -1),
+      use_bias=self.use_bias,
+      rngs=rngs,
+    )
+
+    self.use_per_dim_scale = use_per_dim_scale
+    self.use_rotary_position_embeddings = use_rotary_position_embeddings
+    if self.use_rotary_position_embeddings:
+      self.rotary_position_embedding = RotaryPositionalEmbedding(
+        embedding_dims=self.head_dim,
+      )
+    else:
+      self.rotary_position_embedding = None
+
+    if use_per_dim_scale:
+      self.per_dim_scale = PerDimScale(num_dims=self.head_dim, rngs=rngs)
+    else:
+      self.per_dim_scale = None
+
+  def __call__(
+    self,
+    inputs_q: Array,
+    *,
+    decode_cache: DecodeCache | None = None,
+    patch_mask: Array | None = None,
+    deterministic: bool | None = None,
+    sow_weights: bool = False,
+  ) -> tuple[Float[Array, "b ... o"], DecodeCache | None]:
+    """Applies multi-head dot product attention on the input data."""
+    _, n_patches, input_in_features = inputs_q.shape
+    if input_in_features != self.in_features:
+      raise ValueError(
+        f"Incompatible input dimension, got {input_in_features} "
+        f"but module expects {self.in_features}."
+      )
+    if patch_mask is None:
+      patch_mask = jnp.zeros_like(inputs_q.shape[:-1], dtype=jnp.bool)
+
+    # For query: rope -> ln -> per_dim_scale
+    query = self.query(inputs_q)
+    key = self.key(inputs_q)
+    value = self.value(inputs_q)
+
+    if decode_cache is None:
+      num_masked = jnp.sum(patch_mask.astype(jnp.int32), axis=-1, keepdims=False)
+      next_index = jnp.zeros_like(num_masked, dtype=jnp.int32)
+    else:
+      num_masked = (
+        jnp.sum(patch_mask.astype(jnp.int32), axis=-1, keepdims=False)
+        + decode_cache.num_masked
+      )
+      next_index = decode_cache.next_index
+
+    if self.use_rotary_position_embeddings:
+      position = (
+        jnp.arange(n_patches, dtype=jnp.int32)[None, :]
+        + next_index[:, None]
+        - num_masked[:, None]
+      )
+      query = self.rotary_position_embedding(query, position)
+      key = self.rotary_position_embedding(key, position)
+    if self.query_ln is not None:
+      query = self.query_ln(query)
+    if self.key_ln is not None:
+      key = self.key_ln(key)
+    if self.use_per_dim_scale:
+      query = self.per_dim_scale(query)
+
+    if decode_cache is not None:
+      # Cached decoding.
+      _, decode_cache_size, _, _ = decode_cache.value.shape
+      zero = jnp.array(0, dtype=lax.dtype(next_index.dtype))
+      start_indices = (zero, next_index[0], zero, zero)
+      key = lax.dynamic_update_slice(decode_cache.key, key, start_indices)
+      value = lax.dynamic_update_slice(decode_cache.value, value, start_indices)
+      decode_cache.key = key
+      decode_cache.value = value
+      decode_cache.next_index = next_index + n_patches
+      decode_cache.num_masked = num_masked
+      attn_mask = make_attn_mask(
+        query_length=n_patches,
+        num_all_masked_kv=num_masked,
+        query_index_offset=next_index,
+        kv_length=decode_cache_size,
+      )
+    else:
+      # Training
+      attn_mask = make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked)
+
+    # apply attention
+    x = self.attention_fn(
+      query * jnp.sqrt(self.head_dim),
+      key,
+      value,
+      mask=attn_mask,
+      deterministic=deterministic,
+      module=self if sow_weights else None,
+    )
+    # back to the original inputs dimensions
+    out = self.out(x)
+    return out, decode_cache
+
+
+class Transformer(nnx.Module):
+  """Classic Transformer used in TimesFM."""
+
+  def __init__(self, config: TransformerConfig, *, rngs=nnx.Rngs(42)):
+    self.config = config
+
+    if config.attention_norm == "rms":
+      self.pre_attn_ln = RMSNorm(num_features=config.model_dims, rngs=rngs)
+      self.post_attn_ln = RMSNorm(num_features=config.model_dims, rngs=rngs)
+    else:
+      raise ValueError(f"Layer norm: {config.attention_norm} not supported.")
+
+    self.attn = MultiHeadAttention(
+      num_heads=config.num_heads,
+      in_features=config.model_dims,
+      use_per_dim_scale=True,
+      use_rotary_position_embeddings=config.use_rotary_position_embeddings,
+      qk_norm=config.qk_norm,
+      rngs=rngs,
+    )
+
+    if config.feedforward_norm == "rms":
+      self.pre_ff_ln = RMSNorm(num_features=config.model_dims, rngs=rngs)
+      self.post_ff_ln = RMSNorm(num_features=config.model_dims, rngs=rngs)
+    else:
+      raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.")
+    self.ff0 = nnx.Linear(
+      in_features=config.model_dims,
+      out_features=config.hidden_dims,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+    self.ff1 = nnx.Linear(
+      in_features=config.hidden_dims,
+      out_features=config.model_dims,
+      use_bias=config.use_bias,
+      rngs=rngs,
+    )
+    if config.ff_activation == "relu":
+      self.activation = jax.nn.relu
+    elif config.ff_activation == "swish":
+      self.activation = jax.nn.swish
+    elif config.ff_activation == "none":
+      self.activation = lambda x: x
+    else:
+      raise ValueError(f"Activation: {config.ff_activation} not supported.")
+
+  def __call__(
+    self,
+    input_embeddings: Float[Array, "b n d"],
+    patch_mask: Bool[Array, "b n"],
+    decode_cache: DecodeCache | None = None,
+  ) -> tuple[Float[Array, "b n d"], DecodeCache | None]:
+    attn_output, decode_cache = self.attn(
+      inputs_q=self.pre_attn_ln(input_embeddings),
+      decode_cache=decode_cache,
+      patch_mask=patch_mask,
+      sow_weights=False,
+      deterministic=True,
+    )
+    attn_output = self.post_attn_ln(attn_output) + input_embeddings
+    output_embeddings = (
+      self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
+      + attn_output
+    )
+    return output_embeddings, decode_cache

src/timesfm/flax/util.py

@@ -0,0 +1,107 @@
+# Copyright 2025 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.
+
+"""Flax utility functions for TimesFM layers."""
+
+import dataclasses
+import functools
+import jax
+import jax.numpy as jnp
+import jaxtyping
+
+Float = jaxtyping.Float
+Array = jaxtyping.Array
+Bool = jaxtyping.Bool
+Integer = jaxtyping.Integer
+
+_TOLERANCE = 1e-6
+
+
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=False)
+class DecodeCache:
+  """Cache for decoding."""
+
+  next_index: Integer[Array, "b"]
+  num_masked: Integer[Array, "b"]
+  key: Float[Array, "b n h d"]
+  value: Float[Array, "b n h d"]
+
+
+@jax.jit
+def update_running_stats(
+  n: Float[Array, "b"],
+  mu: Float[Array, "b"],
+  sigma: Float[Array, "b"],
+  x: Float[Array, "b p"],
+  mask: Bool[Array, "b p"],
+) -> tuple[
+  tuple[Float[Array, "b"], Float[Array, "b"], Float[Array, "b"]],
+  tuple[Float[Array, "b"], Float[Array, "b"], Float[Array, "b"]],
+]:
+  """Updates the running stats."""
+  is_legit = jnp.logical_not(mask)
+  inc_n = jnp.sum(is_legit.astype(jnp.float32), axis=-1, keepdims=False)
+  inc_mu = jnp.where(
+    inc_n == 0, 0.0, jnp.mean(x, axis=-1, keepdims=False, where=is_legit)
+  )
+  inc_sigma = jnp.where(
+    inc_n == 0, 0.0, jnp.std(x, axis=-1, keepdims=False, where=is_legit)
+  )
+  new_n = n + inc_n
+  new_mu = jnp.where(new_n == 0, 0.0, (n * mu + inc_mu * inc_n) / new_n)
+  new_sigma = jnp.sqrt(
+    jnp.where(
+      new_n == 0,
+      0.0,
+      (
+        n * sigma * sigma
+        + inc_n * inc_sigma * inc_sigma
+        + n * (mu - new_mu) * (mu - new_mu)
+        + inc_n * (inc_mu - new_mu) * (inc_mu - new_mu)
+      )
+      / new_n,
+    )
+  )
+  return (w := (new_n, new_mu, new_sigma), w)
+
+
+def scan_along_axis(f, init, xs, axis: int, **kwargs):
+  """Scans along an axis."""
+  moved_xs = jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, axis, 0), xs)
+  carry, moved_ys = jax.lax.scan(f, init, moved_xs, **kwargs)
+  return (
+    carry,
+    jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), moved_ys),
+  )
+
+
+@functools.partial(jax.jit, static_argnames=("reverse",))
+def revin(
+  x: Float[Array, "b ..."],
+  mu: Float[Array, "b ..."],
+  sigma: Float[Array, "b ..."],
+  reverse: bool = False,
+):
+  """Reversible per-instance normalization."""
+  if len(mu.shape) == len(x.shape) - 1:
+    mu = mu[..., None]
+    sigma = sigma[..., None]
+  elif len(mu.shape) == len(x.shape) - 2:
+    mu = mu[..., None, None]
+    sigma = sigma[..., None, None]
+  if reverse:
+    return x * sigma + mu
+  else:
+    return (x - mu) / jnp.where(sigma < _TOLERANCE, 1.0, sigma)

src/timesfm/timesfm_2p5/timesfm_2p5_flax.py

@@ -0,0 +1,583 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""TimesFM models in Flax."""
+
+import dataclasses
+import functools
+import gc
+import logging
+import math
+import os
+from pathlib import Path
+from typing import Any, Callable, Dict
+
+import einshape
+from flax import nnx
+import huggingface_hub
+import jax
+import jax.numpy as jnp
+import jaxtyping
+import numpy as np
+import orbax.checkpoint as ocp
+
+from .. import configs
+from ..flax import dense, transformer, util
+from . import timesfm_2p5_base
+
+jax_einshape = einshape.jax_einshape
+scan = util.scan_along_axis
+revin = util.revin
+
+Float = jaxtyping.Float
+Bool = jaxtyping.Bool
+Array = jaxtyping.Array
+
+
+def try_gc():
+  for d in jax.local_devices():
+    stats = d.memory_stats()
+    if stats is None:
+      return
+    if stats["bytes_in_use"] / stats["bytes_limit"] > 0.75:
+      gc.collect()
+      break
+
+
+@nnx.vmap(in_axes=(None, 0), out_axes=0)
+def _create_stacked_transformers(
+  config: configs.StackedTransformersConfig, key: jax.Array
+):
+  return transformer.Transformer(config.transformer, rngs=nnx.Rngs(key))
+
+
+def _scan_along_axis(f, init, xs, axis: int, **kwargs):
+  """Scans along an axis."""
+  moved_xs = jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, axis, 0), xs)
+  carry, moved_ys = jax.lax.scan(f, init, moved_xs, **kwargs)
+  return (
+    carry,
+    jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), moved_ys),
+  )
+
+
+@nnx.scan(in_axes=(0, nnx.Carry, None, 0), out_axes=(nnx.Carry, 0))
+def _apply_stacked_transformers(
+  model: transformer.Transformer,
+  x: Float[Array, "b n d"],
+  m: Float[Array, "b n"],
+  decode_cache: util.DecodeCache | None = None,
+) -> Float[Array, "b n d"]:
+  return model(x, m, decode_cache=decode_cache)
+
+
+class TimesFM_2p5_200M_flax_module(nnx.Module):  # pylint: disable=invalid-name
+  """TimesFM 2.5 with 200M parameters."""
+
+  config = timesfm_2p5_base.TimesFM_2p5_200M_Definition()
+  decode_index: int = 5
+  compiled_decode: Callable[..., Any] | None = None
+  backend: str = ""
+  context: int = 0
+  horizon: int = 0
+  per_core_batch_size: int = 0
+
+  def __init__(self):
+    super().__init__()
+    self.backend = jax.devices()[0].platform
+    self.num_devices = len(jax.devices(self.backend))
+
+    # Names constants.
+    self.p = self.config.input_patch_len  # 32
+    self.o = self.config.output_patch_len  # 128
+    self.os = self.config.output_quantile_len  # 1024
+    self.m = self.o // self.p  # 4
+    self.x = self.config.stacked_transformers.num_layers  # 20
+    self.h = self.config.stacked_transformers.transformer.num_heads  # 16
+    self.md = self.config.stacked_transformers.transformer.model_dims  # 1280
+    self.hd = self.md // self.h  # 80
+    self.q = len(self.config.quantiles) + 1  # 10
+    self.aridx = self.config.decode_index  # 5
+
+    # Layers.
+    self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
+    self.stacked_xf = _create_stacked_transformers(
+      self.config.stacked_transformers,
+      jax.random.split(jax.random.key(42), self.x),
+    )
+    self.output_projection_point = dense.ResidualBlock(
+      self.config.output_projection_point
+    )
+    self.output_projection_quantiles = dense.ResidualBlock(
+      self.config.output_projection_quantiles
+    )
+
+  def __call__(
+    self,
+    inputs: Float[Array, "b n p"],
+    masks: Bool[Array, "b n p"],
+    decode_cache: util.DecodeCache | None = None,
+  ):
+    tokenizer_inputs = jnp.concatenate([inputs, masks.astype(inputs.dtype)], axis=-1)
+    input_embeddings = self.tokenizer(tokenizer_inputs)
+    if decode_cache is None:
+      decode_cache = [None] * self.x
+    output_embeddings, decode_cache = _apply_stacked_transformers(
+      self.stacked_xf, input_embeddings, masks[..., -1], decode_cache
+    )
+    output_ts = self.output_projection_point(output_embeddings)
+    output_quantile_spread = self.output_projection_quantiles(output_embeddings)
+    return (
+      input_embeddings,
+      output_embeddings,
+      output_ts,
+      output_quantile_spread,
+    ), decode_cache
+
+  @nnx.jit(static_argnames=("horizon",))
+  def decode(self, horizon: int, inputs, masks):
+    batch_size, context = inputs.shape[0], inputs.shape[1]
+    num_decode_steps = (horizon - 1) // self.o
+    num_input_patches = context // self.p
+    decode_cache_size = num_input_patches + num_decode_steps * self.m
+
+    # Prefill
+    patched_inputs = jax_einshape("b(np)->bnp", inputs, b=batch_size, p=self.p)
+    patched_masks = jax_einshape("b(np)->bnp", masks, b=batch_size, p=self.p)
+    (last_n, last_mu, last_sigma), (_, context_mu, context_sigma) = scan(
+      lambda carry, xs: util.update_running_stats(*carry, *xs),
+      init=(zero := jnp.zeros(shape=(batch_size)), zero, zero),
+      xs=(patched_inputs, patched_masks),
+      axis=1,
+    )
+    decode_cache = util.DecodeCache(
+      next_index=jnp.zeros(shape=(self.x, batch_size), dtype=jnp.int32),
+      num_masked=jnp.zeros(shape=(self.x, batch_size), dtype=jnp.int32),
+      key=jnp.zeros(shape=(self.x, batch_size, decode_cache_size, self.h, self.hd)),
+      value=jnp.zeros(shape=(self.x, batch_size, decode_cache_size, self.h, self.hd)),
+    )
+    normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False)
+    normed_inputs = jnp.where(patched_masks, 0.0, normed_inputs)
+    (_, _, normed_outputs, normed_quantile_spread), decode_cache = self(
+      normed_inputs, patched_masks, decode_cache
+    )
+    renormed_outputs = jax_einshape(
+      "bn(oq)->bnoq",
+      revin(normed_outputs, context_mu, context_sigma, reverse=True),
+      o=self.o,
+      q=self.q,
+    )
+    renormed_quantile_spread = jax_einshape(
+      "bn(oq)->bnoq",
+      revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
+      o=self.os,
+      q=self.q,
+    )[:, -1, ...]
+
+    # Autogressive decode
+    @nnx.scan(in_axes=(None, nnx.Carry, 0), out_axes=(nnx.Carry, 1))
+    def _ar_decode(module, carry, unused_iter):
+      last_renormed_output, (last_n, last_mu, last_sigma), decode_cache = carry
+      new_patched_input = jax_einshape(
+        "b(mp)->bmp", last_renormed_output, m=module.m, p=module.p
+      )
+      new_mask = jnp.zeros_like(new_patched_input, dtype=jnp.bool)
+      carry_stats, (_, new_mu, new_sigma) = scan(
+        lambda carry, xs: util.update_running_stats(*carry, *xs),
+        init=(last_n, last_mu, last_sigma),
+        xs=(new_patched_input, new_mask),
+        axis=1,
+      )
+      new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False)
+      (_, _, new_normed_output, _), decode_cache = module(
+        new_normed_input, new_mask, decode_cache
+      )
+      new_renormed_output = jax_einshape(
+        "bm(oq)->bmoq",
+        revin(new_normed_output, new_mu, new_sigma, reverse=True),
+        o=module.o,
+        q=module.q,
+      )[..., -1, :, :]
+
+      return (
+        (
+          new_renormed_output[..., module.decode_index],
+          carry_stats,
+          decode_cache,
+        ),
+        new_renormed_output,
+      )
+
+    if num_decode_steps > 0:
+      _, ar_renormed_outputs = _ar_decode(
+        self,
+        (
+          renormed_outputs[..., -1, :, self.decode_index],
+          (last_n, last_mu, last_sigma),
+          decode_cache,
+        ),
+        jnp.arange(num_decode_steps),
+      )
+    else:
+      ar_renormed_outputs = None
+
+    return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs
+
+  def compile(
+    self,
+    context: int,
+    horizon: int,
+    per_core_batch_size: int = 1,
+  ):
+    if context % self.p != 0:
+      logging.info(
+        "When compiling, context needs to be multiple of the patch size %d."
+        " Modifying context to %d.",
+        self.p,
+        context := math.ceil(context / self.p) * self.p,
+      )
+    if horizon % self.o != 0:
+      logging.info(
+        "When compiling, horizon needs to be multiple of the output patch"
+        " size %d. Modifying horizon to %d.",
+        self.o,
+        horizon := math.ceil(horizon / self.o) * self.o,
+      )
+
+    self.context = context
+    self.horizon = horizon
+    self.per_core_batch_size = per_core_batch_size
+
+    @nnx.pmap(
+      in_axes=(None, None, 0, 0),
+      out_axes=(0, 0, 0),
+      devices=jax.devices(self.backend),
+      axis_size=self.num_devices,
+      static_broadcasted_argnums=(1,),
+      axis_name="global_batch",
+    )
+    def compiled_decode_kernel(model, horizon, inputs, masks):
+      return model.decode(horizon, inputs, masks)
+
+    self.compiled_decode = functools.partial(compiled_decode_kernel, self)
+
+
+def _flip_quantile_fn(x):
+  return jnp.concatenate([x[..., :1], jnp.flip(x[..., 1:], axis=-1)], axis=-1)
+
+
+@functools.partial(
+  jax.jit,
+  donate_argnums=(0, 1, 2),
+)
+def _force_flip_invariance_fn(
+  flipped_pf_outputs,
+  flipped_quantile_spreads,
+  flipped_ar_outputs,
+):
+  """Forces flip invariance."""
+  flipped_pf_outputs = _flip_quantile_fn(flipped_pf_outputs)
+  flipped_pf_outputs = jax_einshape("tb...->(tb)...", flipped_pf_outputs)
+  flipped_quantile_spreads = _flip_quantile_fn(flipped_quantile_spreads)
+  flipped_quantile_spreads = jax_einshape("tb...->(tb)...", flipped_quantile_spreads)
+  to_concat = [flipped_pf_outputs[:, -1, ...]]
+  if flipped_ar_outputs is not None:
+    flipped_ar_outputs = _flip_quantile_fn(flipped_ar_outputs)
+    flipped_ar_outputs = jax_einshape("tbno...->(tb)(no)...", flipped_ar_outputs)
+    to_concat.append(flipped_ar_outputs)
+  flipped_full_forecast = jnp.concatenate(to_concat, axis=1)
+
+  return flipped_quantile_spreads, flipped_pf_outputs, flipped_full_forecast
+
+
+@functools.partial(
+  jax.jit,
+  static_argnames=("max_horizon",),
+  donate_argnums=(0,),
+)
+def _use_continuous_quantile_head_fn(full_forecast, quantile_spreads, max_horizon):
+  """Uses continuous quantile head."""
+  to_stack = [full_forecast[..., :max_horizon, 0]]
+  for quantile_index in [1, 2, 3, 4]:
+    to_stack.append(
+      quantile_spreads[:, :max_horizon, quantile_index]
+      - quantile_spreads[:, :max_horizon, 5]
+      + full_forecast[:, :max_horizon, 5]
+    )
+  to_stack.append(full_forecast[..., :max_horizon, 5])
+  for quantile_index in [6, 7, 8, 9]:
+    to_stack.append(
+      quantile_spreads[:, :max_horizon, quantile_index]
+      - quantile_spreads[:, :max_horizon, 5]
+      + full_forecast[:, :max_horizon, 5]
+    )
+  return jnp.stack(to_stack, axis=-1)
+
+
+@functools.partial(jax.jit, donate_argnums=(0,))
+def _fix_quantile_crossing_fn(full_forecast):
+  """Fixes quantile crossing."""
+  lower_quantiles = _scan_along_axis(
+    lambda carry, x: (w := jnp.minimum(carry, x), w),
+    init=full_forecast[..., 5],
+    xs=full_forecast[..., 1:5],
+    axis=-1,
+    reverse=True,
+  )[1]
+  upper_quantiles = _scan_along_axis(
+    lambda carry, x: (w := jnp.maximum(carry, x), w),
+    init=full_forecast[..., 5],
+    xs=full_forecast[..., 6:10],
+    axis=-1,
+    reverse=False,
+  )[1]
+  return jnp.concatenate(
+    [
+      full_forecast[..., :1],
+      lower_quantiles,
+      full_forecast[..., 5:6],
+      upper_quantiles,
+    ],
+    axis=-1,
+  )
+
+
+@functools.partial(jax.jit, static_argnames=("fc",), donate_argnums=(1, 2))
+def _before_model_decode(fc, inputs, masks):
+  """All Jax steps before model decode call."""
+  if fc.infer_is_positive:
+    is_positive = jnp.all(inputs >= 0, axis=-1, keepdims=True)
+  else:
+    is_positive = None
+
+  if fc.normalize_inputs:
+    mu = jnp.mean(inputs, axis=-1, keepdims=True)
+    sigma = jnp.std(inputs, axis=-1, keepdims=True)
+    inputs = revin(inputs, mu, sigma, reverse=False)
+  else:
+    mu, sigma = None, None
+
+  inputs = jax_einshape("(tb)...->tb...", inputs, b=fc.per_core_batch_size)
+  masks = jax_einshape("(tb)...->tb...", masks, b=fc.per_core_batch_size)
+
+  return inputs, masks, is_positive, mu, sigma
+
+
+@functools.partial(
+  jax.jit,
+  static_argnames=(
+    "fc",
+    "p",
+  ),
+  donate_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9),
+)
+def _after_model_decode(
+  fc,
+  pf_outputs,
+  quantile_spreads,
+  ar_outputs,
+  flipped_pf_outputs,
+  flipped_quantile_spreads,
+  flipped_ar_outputs,
+  is_positive,
+  mu,
+  sigma,
+  p,
+):
+  """All Jax steps after model decode call."""
+  # t: num_devices, b: per_core_batch_size
+  pf_outputs = jax_einshape("tb...->(tb)...", pf_outputs)
+  quantile_spreads = jax_einshape("tb...->(tb)...", quantile_spreads)
+  to_concat = [pf_outputs[:, -1, ...]]
+  if ar_outputs is not None:
+    ar_outputs = jax_einshape("tbno...->(tb)(no)...", ar_outputs)
+    to_concat.append(ar_outputs)
+  full_forecast = jnp.concatenate(to_concat, axis=1)
+
+  if fc.force_flip_invariance:
+    (
+      flipped_quantile_spreads,
+      flipped_pf_outputs,
+      flipped_full_forecast,
+    ) = _force_flip_invariance_fn(
+      flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs
+    )
+    quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2
+    pf_outputs = (pf_outputs - flipped_pf_outputs) / 2
+    full_forecast = (full_forecast - flipped_full_forecast) / 2
+
+  if fc.use_continuous_quantile_head:
+    full_forecast = _use_continuous_quantile_head_fn(
+      full_forecast, quantile_spreads, fc.max_horizon
+    )
+
+  if fc.return_backcast:
+    full_backcast = jax_einshape("...npq->...(np)q", pf_outputs[:, :-1, :p, :])
+    full_forecast = jnp.concatenate([full_backcast, full_forecast], axis=1)
+
+  if fc.fix_quantile_crossing:
+    full_forecast = _fix_quantile_crossing_fn(full_forecast)
+
+  if fc.normalize_inputs:
+    full_forecast = revin(full_forecast, mu, sigma, reverse=True)
+
+  if is_positive is not None:
+    full_forecast = jnp.where(
+      is_positive[..., None],
+      jnp.maximum(full_forecast, jnp.zeros_like(full_forecast)),
+      full_forecast,
+    )
+
+  return full_forecast
+
+
+class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
+  """Flax implementation of TimesFM 2.5 with 200M parameters."""
+
+  model: nnx.Module = TimesFM_2p5_200M_flax_module()
+
+  @classmethod
+  def from_pretrained(
+    cls,
+    model_id: str = "google/timesfm-2.5-200m-flax",
+    *,
+    revision: str | None = None,
+    cache_dir: str | Path | None = None,
+    force_download: bool = False,
+    proxies: Dict | None = None,
+    resume_download: bool | None = None,
+    local_files_only: bool | None = None,
+    token: str | None = None,
+    **model_kwargs,
+  ):
+    """Loads a Flax TimesFM model."""
+
+    # Create an instance of the model wrapper class.
+    instance = cls(**model_kwargs)
+
+    # Determine the path to the model weights.
+    model_file_path = ""
+    if os.path.isdir(model_id):
+      logging.info("Loading checkpoint from local directory: %s", model_id)
+      model_file_path = model_id
+    else:
+      logging.info("Downloading checkpoint from Hugging Face repo %s", model_id)
+      model_file_path = huggingface_hub.snapshot_download(
+        repo_id=model_id,
+        revision=revision,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        proxies=proxies,
+        resume_download=resume_download,
+        token=token,
+        local_files_only=local_files_only,
+      )
+      logging.info("Loading checkpoint from: %s", model_file_path)
+
+    checkpointer = ocp.StandardCheckpointer()
+    graph, state = nnx.split(instance.model)
+    state = checkpointer.restore(model_file_path, state)
+    instance.model = nnx.merge(graph, state)
+    return instance
+
+  def compile(self, forecast_config: configs.ForecastConfig, **kwargs):
+    # Acrobym used during validation.
+    fc = forecast_config
+    if fc.max_context % self.model.p != 0:
+      logging.info(
+        "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,
+      )
+      fc = dataclasses.replace(fc, max_context=new_context)
+    if fc.max_horizon % self.model.o != 0:
+      logging.info(
+        "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,
+      )
+      fc = dataclasses.replace(fc, max_horizon=new_horizon)
+    if fc.max_context + fc.max_horizon > self.model.config.context_limit:
+      raise ValueError(
+        "Context + horizon must be less than the context limit."
+        f" {fc.max_context} + {fc.max_horizon} >"
+        f" {self.model.config.context_limit}."
+      )
+    if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os):
+      raise ValueError(
+        f"Continuous quantile head is not supported for horizons > {self.model.os}."
+      )
+
+    self.forecast_config = forecast_config
+    self.model.compile(
+      context=self.forecast_config.max_context,
+      horizon=self.forecast_config.max_horizon,
+      per_core_batch_size=fc.per_core_batch_size,
+    )
+    self.per_core_batch_size = self.forecast_config.per_core_batch_size
+    self.num_devices = self.model.num_devices
+    self.global_batch_size = (
+      self.forecast_config.per_core_batch_size * self.model.num_devices
+    )
+
+    def compiled_decode_kernel(fc, horizon, inputs, masks):
+      inputs = jnp.array(inputs, dtype=jnp.float32)
+      masks = jnp.array(masks, dtype=jnp.bool)
+      if horizon > fc.max_horizon:
+        raise ValueError(
+          f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}."
+        )
+      to_trim = fc.max_horizon - horizon
+
+      inputs, masks, is_positive, mu, sigma = _before_model_decode(fc, inputs, masks)
+
+      pf_outputs, quantile_spreads, ar_outputs = self.model.compiled_decode(
+        fc.max_horizon, inputs, masks
+      )
+      if fc.force_flip_invariance:
+        flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
+          self.model.compiled_decode(fc.max_horizon, -inputs, masks)
+        )
+      else:
+        flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
+          None,
+          None,
+          None,
+        )
+
+      full_forecast = _after_model_decode(
+        fc,
+        pf_outputs,
+        quantile_spreads,
+        ar_outputs,
+        flipped_pf_outputs,
+        flipped_quantile_spreads,
+        flipped_ar_outputs,
+        is_positive,
+        mu,
+        sigma,
+        self.model.p,
+      )
+      full_forecast_np = np.array(full_forecast)
+      del full_forecast
+      try_gc()
+      if to_trim > 0:
+        full_forecast_np = full_forecast_np[..., :-to_trim, :]
+      return full_forecast_np[..., 5], full_forecast_np
+
+    self.compiled_decode = functools.partial(
+      compiled_decode_kernel, self.forecast_config
+    )

src/timesfm/timesfm_2p5/timesfm_2p5_torch.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """TimesFM models."""
 
+import dataclasses
 import logging
 import math
 import os
@@ -54,12 +55,10 @@ class TimesFM_2p5_200M_torch_module(nn.Module):
 
     # Layers.
     self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
-    self.stacked_xf = nn.ModuleList(
-      [
-        transformer.Transformer(self.config.stacked_transformers.transformer)
-        for _ in range(self.x)
-      ]
-    )
+    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
     )
@@ -271,7 +270,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
   def _from_pretrained(
     cls,
     *,
-    model_id: str,
+    model_id: str = "google/timesfm-2.5-200m-pytorch",
     revision: Optional[str],
     cache_dir: Optional[Union[str, Path]],
     force_download: bool,
@@ -349,7 +348,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
         self.model.p,
         new_context := math.ceil(fc.max_context / self.model.p) * self.model.p,
       )
-      fc.max_context = new_context
+      fc = dataclasses.replace(fc, max_context=new_context)
     if fc.max_horizon % self.model.o != 0:
       logging.info(
         "When compiling, max horizon needs to be multiple of the output patch"
@@ -357,7 +356,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
         self.model.o,
         new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o,
       )
-      fc.max_horizon = new_horizon
+      fc = dataclasses.replace(fc, max_horizon=new_horizon)
     if fc.max_context + fc.max_horizon > self.model.config.context_limit:
       raise ValueError(
         "Context + horizon must be less than the context limit."