freedak/timesfm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

minor changes

Browse Source
This commit is contained in:
siriuz42
2025-10-07 20:41:38 +00:00
parent 25bb4f20e8
commit ff2cefaa26
7 changed files with 366 additions and 389 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/timesfm/flax/dense.py
+23 -23
View File
@@ -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
src/timesfm/flax/normalization.py
+8 -15
View File
@@ -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)
src/timesfm/flax/transformer.py
+92 -101
View File
@@ -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
src/timesfm/flax/util.py
+25 -26
View File
@@ -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:
src/timesfm/timesfm_2p5/timesfm_2p5_flax.py
+213 -217
View File
@@ -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
)
src/timesfm/timesfm_2p5/timesfm_2p5_torch.py
+5 -7
View File
@@ -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,
where
View File