diff --git a/pyproject.toml b/pyproject.toml index 3cb0346..c5344be 100644 --- a/pyproject.toml +++ b/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 diff --git a/src/timesfm/configs.py b/src/timesfm/configs.py index 104ea42..88a63af 100644 --- a/src/timesfm/configs.py +++ b/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. diff --git a/src/timesfm/flax/dense.py b/src/timesfm/flax/dense.py index 7728d36..7de49b9 100644 --- a/src/timesfm/flax/dense.py +++ b/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) diff --git a/src/timesfm/flax/normalization.py b/src/timesfm/flax/normalization.py index 6f6e0be..07349c8 100644 --- a/src/timesfm/flax/normalization.py +++ b/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"]: diff --git a/src/timesfm/flax/transformer.py b/src/timesfm/flax/transformer.py index fb56d4b..add1d85 100644 --- a/src/timesfm/flax/transformer.py +++ b/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"], diff --git a/src/timesfm/flax/util.py b/src/timesfm/flax/util.py index d98ef5c..de4d3c7 100644 --- a/src/timesfm/flax/util.py +++ b/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"], diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py index 87c09d9..48a61ad 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py +++ 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): diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 897dbe7..260cbd3 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/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."