diff --git a/README.md b/README.md index 3479b00..9db8441 100644 --- a/README.md +++ b/README.md @@ -58,13 +58,10 @@ will be under construction over the next few weeks to # Activate the environment source .venv/bin/activate - # Install the dependencies - uv pip install -r requirements.txt - ``` - -3. (Optional) If you plan to modify the source code, install the package in editable mode: - ```shell - uv pip install -e . + # Install the package in editable mode with torch + uv pip install -e .[torch] + # Or with flax + uv pip install -e .[flax] ``` ### Code Example diff --git a/pyproject.toml b/pyproject.toml index ceb12f4..715240b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,8 +16,20 @@ dependencies = [ "numpy>=1.26.4", "huggingface_hub[cli]>=0.23.0", "safetensors>=0.5.3", +] + +[project.optional-dependencies] +torch = [ "torch>=2.0.0", ] +flax = [ + "flax>=0.8.2", + "optax>=0.2.2", + "einshape>=0.8.0", + "orbax-checkpoint>=0.5.15", + "jaxtyping>=0.2.29", + "jax>=0.4.26", +] [tool.ruff] line-length = 88 diff --git a/src/timesfm/__init__.py b/src/timesfm/__init__.py index 85177a2..361a36d 100644 --- a/src/timesfm/__init__.py +++ b/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 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/__init__.py b/src/timesfm/flax/__init__.py new file mode 100644 index 0000000..0a2669d --- /dev/null +++ b/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. diff --git a/src/timesfm/flax/dense.py b/src/timesfm/flax/dense.py new file mode 100644 index 0000000..6431912 --- /dev/null +++ b/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 diff --git a/src/timesfm/flax/normalization.py b/src/timesfm/flax/normalization.py new file mode 100644 index 0000000..390e44f --- /dev/null +++ b/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 diff --git a/src/timesfm/flax/transformer.py b/src/timesfm/flax/transformer.py new file mode 100644 index 0000000..6e61967 --- /dev/null +++ b/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 diff --git a/src/timesfm/flax/util.py b/src/timesfm/flax/util.py new file mode 100644 index 0000000..ec70d72 --- /dev/null +++ b/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) diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py new file mode 100644 index 0000000..acd354b --- /dev/null +++ b/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 + ) diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py index 897dbe7..c50be89 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 @@ -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."