(1) Make config hashable. (2) Flax

pyproject.toml

@@ -24,13 +24,12 @@ 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 }
-typeguard = { version = ">=4.3.0", optional = true }
 jax = { version = ">=0.4.26", optional = true }
 
 
 [tool.poetry.extras]
 torch = ["torch"]
-flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "typeguard", "jax"]
+flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "jax"]
 
 [tool.ruff]
 line-length = 88

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/dense.py

@@ -18,7 +18,6 @@ from flax import nnx
 import jax
 import jax.numpy as jnp
 import jaxtyping
-import typeguard
 
 from .. import configs
 
@@ -64,7 +63,6 @@ class ResidualBlock(nnx.Module):
     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))
@@ -99,7 +97,6 @@ class RandomFourierFeatures(nnx.Module):
         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)

src/timesfm/flax/normalization.py

@@ -18,7 +18,6 @@ from flax import nnx
 import jax
 import jax.numpy as jnp
 import jaxtyping
-import typeguard
 
 Array = jaxtyping.Array
 Bool = jaxtyping.Bool
@@ -44,7 +43,6 @@ class RMSNorm(nnx.Module):
     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"]:
@@ -69,7 +67,6 @@ class LayerNorm(nnx.Module):
     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"]:

src/timesfm/flax/transformer.py

@@ -23,7 +23,6 @@ import jax
 from jax import lax
 import jax.numpy as jnp
 import jaxtyping
-import typeguard
 
 from .. import configs
 from . import normalization, util
@@ -127,7 +126,6 @@ class PerDimScale(nnx.Module):
     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
@@ -344,7 +342,6 @@ class Transformer(nnx.Module):
     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"],

src/timesfm/flax/util.py

@@ -40,7 +40,6 @@ class DecodeCache:
   value: Float[Array, "b n h d"]
 
 
-@jaxtyping.jaxtyped(typechecker=typeguard.typechecked)
 @jax.jit
 def update_running_stats(
     n: Float[Array, "b"],

src/timesfm/timesfm_2p5/timesfm_2p5_flax.py

@@ -46,6 +46,8 @@ 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
@@ -458,9 +460,10 @@ 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(
-      self,
+      cls,
       *,
       path: str | None = None,
       hf_repo_id: str | None = "google/timesfm-2.5-200m-flax",
@@ -476,11 +479,13 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
       logging.info("Loading checkpoint from: %s", path)
     else:
       raise ValueError("Either path or hf_repo_id must be provided.")
-
+    
+    instance = cls()
     checkpointer = ocp.StandardCheckpointer()
-    graph, state = nnx.split(self.model)
+    graph, state = nnx.split(instance.model)
     state = checkpointer.restore(path, state)
-    self.model = nnx.merge(graph, state)
+    instance.model = nnx.merge(graph, state)
+    return instance
 
   def compile(self, forecast_config: configs.ForecastConfig, **kwargs):
 

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
@@ -349,7 +350,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 +358,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
+      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."