From ff2cefaa26515cb8d460246301b51930a196b139 Mon Sep 17 00:00:00 2001 From: siriuz42 Date: Tue, 7 Oct 2025 20:41:38 +0000 Subject: [PATCH] minor changes --- src/timesfm/flax/dense.py | 46 +- src/timesfm/flax/normalization.py | 23 +- src/timesfm/flax/transformer.py | 193 ++++----- src/timesfm/flax/util.py | 51 ++- src/timesfm/timesfm_2p5/timesfm_2p5_flax.py | 430 +++++++++---------- src/timesfm/timesfm_2p5/timesfm_2p5_torch.py | 12 +- where | 0 7 files changed, 366 insertions(+), 389 deletions(-) delete mode 100644 where diff --git a/src/timesfm/flax/dense.py b/src/timesfm/flax/dense.py index 7de49b9..6431912 100644 --- a/src/timesfm/flax/dense.py +++ b/src/timesfm/flax/dense.py @@ -37,22 +37,22 @@ class ResidualBlock(nnx.Module): 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, + 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, + 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, + 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 @@ -65,7 +65,7 @@ class ResidualBlock(nnx.Module): def __call__(self, x: Float[Array, "b ... i"]) -> Float[Array, "b ... o"]: return self.output_layer( - self.activation(self.hidden_layer(x)) + self.activation(self.hidden_layer(x)) ) + self.residual_layer(x) @@ -79,22 +79,22 @@ class RandomFourierFeatures(nnx.Module): if config.output_dims % 4 != 0: raise ValueError( - f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0." + 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, + 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, + 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"]: @@ -104,7 +104,7 @@ class RandomFourierFeatures(nnx.Module): 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 + [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 index 07349c8..390e44f 100644 --- a/src/timesfm/flax/normalization.py +++ b/src/timesfm/flax/normalization.py @@ -32,21 +32,18 @@ class RMSNorm(nnx.Module): __data__ = ("scale",) def __init__( - self, - num_features: int, - *, - epsilon: float = 1e-6, - rngs=nnx.Rngs(42), + 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"]: - + 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 @@ -58,18 +55,14 @@ class LayerNorm(nnx.Module): __data__ = ("scale", "bias") - def __init__( - self, num_features: int, *, epsilon: float = 1e-6, rngs=nnx.Rngs(42) - ): + 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"]: + 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) diff --git a/src/timesfm/flax/transformer.py b/src/timesfm/flax/transformer.py index add1d85..6e61967 100644 --- a/src/timesfm/flax/transformer.py +++ b/src/timesfm/flax/transformer.py @@ -40,14 +40,14 @@ DecodeCache = util.DecodeCache @functools.partial( - jax.jit, - static_argnames=("query_length", "kv_length"), + 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, + 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.""" @@ -59,8 +59,8 @@ def make_attn_mask( 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], + q_index >= kv_index, + kv_index >= num_all_masked_kv[:, None, None, None], ) @@ -68,31 +68,30 @@ class RotaryPositionalEmbedding(nnx.Module): """Rotary positional embedding.""" def __init__( - self, - embedding_dims: int, - min_timescale: int = 1, - max_timescale: int = 10000, + 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, + 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." + "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 + self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction ) if position is None: seq_length = inputs.shape[1] @@ -128,9 +127,7 @@ class PerDimScale(nnx.Module): 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) + 1.442695041 / jnp.sqrt(self.num_dims) * jax.nn.softplus(self.per_dim_scale) ) @@ -138,17 +135,17 @@ 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: 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 @@ -162,15 +159,15 @@ class MultiHeadAttention(nnx.Module): 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})." + 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, + 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] @@ -186,18 +183,18 @@ class MultiHeadAttention(nnx.Module): 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, + 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, + embedding_dims=self.head_dim, ) else: self.rotary_position_embedding = None @@ -208,20 +205,20 @@ class MultiHeadAttention(nnx.Module): 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, + 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}." + 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) @@ -232,22 +229,20 @@ class MultiHeadAttention(nnx.Module): value = self.value(inputs_q) if decode_cache is None: - num_masked = jnp.sum( - patch_mask.astype(jnp.int32), axis=-1, keepdims=False - ) + 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 + 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] + 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) @@ -270,25 +265,23 @@ class MultiHeadAttention(nnx.Module): 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, + 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 - ) + 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, + 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) @@ -308,12 +301,12 @@ class Transformer(nnx.Module): 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, + 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": @@ -322,16 +315,16 @@ class Transformer(nnx.Module): 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, + 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, + 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 @@ -343,23 +336,21 @@ class Transformer(nnx.Module): 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, + 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, + 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 + 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 index de4d3c7..ec70d72 100644 --- a/src/timesfm/flax/util.py +++ b/src/timesfm/flax/util.py @@ -19,7 +19,6 @@ import functools import jax import jax.numpy as jnp import jaxtyping -import typeguard Float = jaxtyping.Float Array = jaxtyping.Array @@ -42,38 +41,38 @@ class DecodeCache: @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"], + 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"]], + 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_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) + 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, + 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) @@ -83,17 +82,17 @@ def scan_along_axis(f, init, xs, axis: int, **kwargs): 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), + 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, + 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: diff --git a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py index 48a61ad..aec4079 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_flax.py @@ -19,7 +19,9 @@ import functools import gc import logging import math -from typing import Any, Callable +import os +from pathlib import Path +from typing import Any, Callable, Dict import einshape from flax import nnx @@ -55,7 +57,7 @@ def try_gc(): @nnx.vmap(in_axes=(None, 0), out_axes=0) def _create_stacked_transformers( - config: configs.StackedTransformersConfig, key: jax.Array + config: configs.StackedTransformersConfig, key: jax.Array ): return transformer.Transformer(config.transformer, rngs=nnx.Rngs(key)) @@ -65,17 +67,17 @@ def _scan_along_axis(f, init, xs, axis: int, **kwargs): 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), + 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, + 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) @@ -111,38 +113,36 @@ class TimesFM_2p5_200M_flax_module(nnx.Module): # pylint: disable=invalid-name # 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.config.stacked_transformers, + jax.random.split(jax.random.key(42), self.x), ) self.output_projection_point = dense.ResidualBlock( - self.config.output_projection_point + self.config.output_projection_point ) self.output_projection_quantiles = dense.ResidualBlock( - self.config.output_projection_quantiles + 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, + 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 - ) + 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 + 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, + input_embeddings, + output_embeddings, + output_ts, + output_quantile_spread, ), decode_cache @nnx.jit(static_argnames=("horizon",)) @@ -156,39 +156,33 @@ class TimesFM_2p5_200M_flax_module(nnx.Module): # pylint: disable=invalid-name 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, + 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 + 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 + 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, + "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, + "bn(oq)->bnoq", + revin(normed_quantile_spread, context_mu, context_sigma, reverse=True), + o=self.os, + q=self.q, )[:, -1, ...] # Autogressive decode @@ -196,46 +190,44 @@ class TimesFM_2p5_200M_flax_module(nnx.Module): # pylint: disable=invalid-name 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 + "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 + 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_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, + "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, + ( + 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), + 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 @@ -243,24 +235,24 @@ class TimesFM_2p5_200M_flax_module(nnx.Module): # pylint: disable=invalid-name return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs def compile( - self, - context: int, - horizon: int, - per_core_batch_size: int = 1, + 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, + "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, + "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 @@ -268,12 +260,12 @@ class TimesFM_2p5_200M_flax_module(nnx.Module): # pylint: disable=invalid-name 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", + 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) @@ -286,27 +278,23 @@ def _flip_quantile_fn(x): @functools.partial( - jax.jit, - donate_argnums=(0, 1, 2), + jax.jit, + donate_argnums=(0, 1, 2), ) def _force_flip_invariance_fn( - flipped_pf_outputs, - flipped_quantile_spreads, - flipped_ar_outputs, + 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 - ) + 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 - ) + 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) @@ -314,27 +302,25 @@ def _force_flip_invariance_fn( @functools.partial( - jax.jit, - static_argnames=("max_horizon",), - donate_argnums=(0,), + jax.jit, + static_argnames=("max_horizon",), + donate_argnums=(0,), ) -def _use_continuous_quantile_head_fn( - full_forecast, quantile_spreads, max_horizon -): +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] + 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] + quantile_spreads[:, :max_horizon, quantile_index] + - quantile_spreads[:, :max_horizon, 5] + + full_forecast[:, :max_horizon, 5] ) return jnp.stack(to_stack, axis=-1) @@ -343,27 +329,27 @@ def _use_continuous_quantile_head_fn( 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, + 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, + 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, + [ + full_forecast[..., :1], + lower_quantiles, + full_forecast[..., 5:6], + upper_quantiles, + ], + axis=-1, ) @@ -389,25 +375,25 @@ def _before_model_decode(fc, inputs, masks): @functools.partial( - jax.jit, - static_argnames=( - "fc", - "p", - ), - donate_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9), + 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, + 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 @@ -421,11 +407,11 @@ def _after_model_decode( if fc.force_flip_invariance: ( - flipped_quantile_spreads, - flipped_pf_outputs, - flipped_full_forecast, + flipped_quantile_spreads, + flipped_pf_outputs, + flipped_full_forecast, ) = _force_flip_invariance_fn( - flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs + 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 @@ -433,7 +419,7 @@ def _after_model_decode( if fc.use_continuous_quantile_head: full_forecast = _use_continuous_quantile_head_fn( - full_forecast, quantile_spreads, fc.max_horizon + full_forecast, quantile_spreads, fc.max_horizon ) if fc.return_backcast: @@ -448,9 +434,9 @@ def _after_model_decode( if is_positive is not None: full_forecast = jnp.where( - is_positive[..., None], - jnp.maximum(full_forecast, jnp.zeros_like(full_forecast)), - full_forecast, + is_positive[..., None], + jnp.maximum(full_forecast, jnp.zeros_like(full_forecast)), + full_forecast, ) return full_forecast @@ -460,120 +446,130 @@ 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, - *, - path: str | None = None, - hf_repo_id: str | None = "google/timesfm-2.5-200m-flax", + 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.""" - 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) + + # 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: - raise ValueError("Either path or hf_repo_id must be provided.") - - instance = cls() + 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(path, state) + 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, + "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, + "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}." + "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}." + 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, + 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 + 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}." + 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 - ) + 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 + 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) + self.model.compiled_decode(fc.max_horizon, -inputs, masks) ) else: flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( - None, - None, - None, + 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, + 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 @@ -583,5 +579,5 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5): return full_forecast_np[..., 5], full_forecast_np self.compiled_decode = functools.partial( - compiled_decode_kernel, self.forecast_config + 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 260cbd3..fa57f14 100644 --- a/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py +++ b/src/timesfm/timesfm_2p5/timesfm_2p5_torch.py @@ -55,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 ) @@ -272,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, diff --git a/where b/where deleted file mode 100644 index e69de29..0000000