initial flax push

pyproject.toml

@@ -14,11 +14,18 @@ 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 }
+jax = { version = ">=0.4.26", optional = true }
+optax = { version = ">=0.2.2", optional = true }
+
+[tool.poetry.extras]
+torch = ["torch"]
+flax = ["flax", "jax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "typeguard"]
 
 [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/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,113 @@
+# 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
+import typeguard
+
+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.")
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  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,
+    )
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  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,81 @@
+# 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
+import typeguard
+
+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
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  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
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  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,368 @@
+# 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
+import typeguard
+
+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,)))
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  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.")
+
+  @jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+  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,109 @@
+# 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
+import typeguard
+
+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"]
+
+
+@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
+@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,582 @@
+# 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
+from typing import Any, Callable
+
+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["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()
+
+  def from_pretrained(
+      self,
+      *,
+      path: str | None = None,
+      hf_repo_id: str | None = "google/timesfm-2.5-200m-flax",
+  ):
+    """Loads a Flax TimesFM model."""
+    if path:
+      pass
+    elif hf_repo_id:
+      logging.info(
+          "Downloading checkpoint from HuggingFace repo %s", hf_repo_id
+      )
+      path = huggingface_hub.snapshot_download(hf_repo_id)
+      logging.info("Loading checkpoint from: %s", path)
+    else:
+      raise ValueError("Either path or hf_repo_id must be provided.")
+
+    checkpointer = ocp.StandardCheckpointer()
+    graph, state = nnx.split(self.model)
+    state = checkpointer.restore(path, state)
+    self.model = nnx.merge(graph, state)
+
+  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,
+      )
+      forecast_config = 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(
+          "Continuous quantile head is not supported for horizons >"
+          f" {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(
+            "Horizon must be less than the max horizon."
+            f" {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
+    )