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)): def __init__(self, config: ResidualBlockConfig, *, rngs=nnx.Rngs(42)):
self.config = config self.config = config
self.hidden_layer = nnx.Linear( self.hidden_layer = nnx.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=config.hidden_dims, out_features=config.hidden_dims,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
self.output_layer = nnx.Linear( self.output_layer = nnx.Linear(
in_features=config.hidden_dims, in_features=config.hidden_dims,
out_features=config.output_dims, out_features=config.output_dims,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
self.residual_layer = nnx.Linear( self.residual_layer = nnx.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=config.output_dims, out_features=config.output_dims,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
if config.activation == "relu": if config.activation == "relu":
self.activation = jax.nn.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"]: def __call__(self, x: Float[Array, "b ... i"]) -> Float[Array, "b ... o"]:
return self.output_layer( return self.output_layer(
self.activation(self.hidden_layer(x)) self.activation(self.hidden_layer(x))
) + self.residual_layer(x) ) + self.residual_layer(x)
@@ -79,22 +79,22 @@ class RandomFourierFeatures(nnx.Module):
if config.output_dims % 4 != 0: if config.output_dims % 4 != 0:
raise ValueError( 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 num_projected_features = config.output_dims // 4
self.phase_shifts = nnx.Param(jnp.zeros(shape=(2, num_projected_features))) self.phase_shifts = nnx.Param(jnp.zeros(shape=(2, num_projected_features)))
self.projection_layer = nnx.Linear( self.projection_layer = nnx.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=num_projected_features, out_features=num_projected_features,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
self.residual_layer = nnx.Linear( self.residual_layer = nnx.Linear(
in_features=config.input_dims, in_features=config.input_dims,
out_features=config.output_dims, out_features=config.output_dims,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
def __call__(self, x: Float[Array, "b ... i"]) -> Float[Array, "b ... o"]: 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_1 = jnp.sign(jnp.sin(projected + self.phase_shifts[0, :]))
sq_wave_2 = jnp.sign(jnp.sin(projected + self.phase_shifts[1, :])) sq_wave_2 = jnp.sign(jnp.sin(projected + self.phase_shifts[1, :]))
fourier_features = jnp.concatenate( 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) residual = self.residual_layer(x)
return fourier_features + residual return fourier_features + residual
src/timesfm/flax/normalization.py
+8 -15
View File
@@ -32,21 +32,18 @@ class RMSNorm(nnx.Module):
__data__ = ("scale",) __data__ = ("scale",)
def __init__( def __init__(
self, self,
num_features: int, num_features: int,
*, *,
epsilon: float = 1e-6, epsilon: float = 1e-6,
rngs=nnx.Rngs(42), rngs=nnx.Rngs(42),
): ):
del rngs del rngs
self.scale = nnx.Param(jnp.zeros(shape=(num_features,))) self.scale = nnx.Param(jnp.zeros(shape=(num_features,)))
self.num_features = num_features self.num_features = num_features
self.epsilon = epsilon self.epsilon = epsilon
def __call__( def __call__(self, inputs: Float[Array, "b ... d"]) -> Float[Array, "b ... d"]:
self, inputs: Float[Array, "b ... d"]
) -> Float[Array, "b ... d"]:
var = jnp.mean(jnp.square(inputs), axis=-1, keepdims=True) var = jnp.mean(jnp.square(inputs), axis=-1, keepdims=True)
normed_inputs = inputs * jax.lax.rsqrt(var + self.epsilon) normed_inputs = inputs * jax.lax.rsqrt(var + self.epsilon)
normed_inputs *= self.scale normed_inputs *= self.scale
@@ -58,18 +55,14 @@ class LayerNorm(nnx.Module):
__data__ = ("scale", "bias") __data__ = ("scale", "bias")
def __init__( 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 del rngs
self.scale = nnx.Param(jnp.ones(shape=(num_features,))) self.scale = nnx.Param(jnp.ones(shape=(num_features,)))
self.bias = nnx.Param(jnp.zeros(shape=(num_features,))) self.bias = nnx.Param(jnp.zeros(shape=(num_features,)))
self.num_features = num_features self.num_features = num_features
self.epsilon = epsilon self.epsilon = epsilon
def __call__( def __call__(self, inputs: Float[Array, "b ... d"]) -> Float[Array, "b ... d"]:
self, inputs: Float[Array, "b ... d"]
) -> Float[Array, "b ... d"]:
mean = jnp.mean(inputs, axis=-1, keepdims=True) mean = jnp.mean(inputs, axis=-1, keepdims=True)
var = jnp.mean(jnp.square(inputs - mean), 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 = (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( @functools.partial(
jax.jit, jax.jit,
static_argnames=("query_length", "kv_length"), static_argnames=("query_length", "kv_length"),
) )
def make_attn_mask( def make_attn_mask(
query_length: int, query_length: int,
num_all_masked_kv: Integer[Array, "b"], num_all_masked_kv: Integer[Array, "b"],
query_index_offset: Integer[Array, "b"] | None = None, query_index_offset: Integer[Array, "b"] | None = None,
kv_length: int = 0, kv_length: int = 0,
) -> Bool[Array, "b 1 q n"]: ) -> Bool[Array, "b 1 q n"]:
"""Makes attention mask.""" """Makes attention mask."""
@@ -59,8 +59,8 @@ def make_attn_mask(
q_index += query_index_offset[:, None, None, None] q_index += query_index_offset[:, None, None, None]
kv_index = jnp.arange(kv_length)[None, None, None, :] kv_index = jnp.arange(kv_length)[None, None, None, :]
return jnp.logical_and( return jnp.logical_and(
q_index >= kv_index, q_index >= kv_index,
kv_index >= num_all_masked_kv[:, None, None, None], kv_index >= num_all_masked_kv[:, None, None, None],
) )
@@ -68,31 +68,30 @@ class RotaryPositionalEmbedding(nnx.Module):
"""Rotary positional embedding.""" """Rotary positional embedding."""
def __init__( def __init__(
self, self,
embedding_dims: int, embedding_dims: int,
min_timescale: int = 1, min_timescale: int = 1,
max_timescale: int = 10000, max_timescale: int = 10000,
): ):
self.embedding_dims = embedding_dims self.embedding_dims = embedding_dims
self.min_timescale = min_timescale self.min_timescale = min_timescale
self.max_timescale = max_timescale self.max_timescale = max_timescale
def __call__( def __call__(
self, self,
inputs: Float[Array, "b ... d"], inputs: Float[Array, "b ... d"],
position: Array | None = None, position: Array | None = None,
): ):
"""Generates a JTensor of sinusoids with different frequencies.""" """Generates a JTensor of sinusoids with different frequencies."""
if self.embedding_dims != inputs.shape[-1]: if self.embedding_dims != inputs.shape[-1]:
raise ValueError( raise ValueError(
"The embedding dims of the rotary position embedding" "The embedding dims of the rotary position embedding"
"must match the hidden dimension of the inputs." "must match the hidden dimension of the inputs."
) )
half_embedding_dim = self.embedding_dims // 2 half_embedding_dim = self.embedding_dims // 2
fraction = 2 * jnp.arange(0, half_embedding_dim) / self.embedding_dims fraction = 2 * jnp.arange(0, half_embedding_dim) / self.embedding_dims
timescale = ( timescale = (
self.min_timescale self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction
* (self.max_timescale / self.min_timescale) ** fraction
) )
if position is None: if position is None:
seq_length = inputs.shape[1] 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"]: def __call__(self, x: Float[Array, "b ... d"]) -> Float[Array, "b ... d"]:
return x * ( return x * (
1.442695041 1.442695041 / jnp.sqrt(self.num_dims) * jax.nn.softplus(self.per_dim_scale)
/ jnp.sqrt(self.num_dims)
* jax.nn.softplus(self.per_dim_scale)
) )
@@ -138,17 +135,17 @@ class MultiHeadAttention(nnx.Module):
"""Multi-head attention.""" """Multi-head attention."""
def __init__( def __init__(
self, self,
num_heads: int, num_heads: int,
in_features: int, in_features: int,
*, *,
use_per_dim_scale: bool = True, use_per_dim_scale: bool = True,
use_rotary_position_embeddings: bool = True, use_rotary_position_embeddings: bool = True,
use_bias: bool = False, use_bias: bool = False,
deterministic: bool | None = None, deterministic: bool | None = None,
attention_fn: Callable[..., Array] = nnx.dot_product_attention, attention_fn: Callable[..., Array] = nnx.dot_product_attention,
qk_norm: str = "rms", qk_norm: str = "rms",
rngs=nnx.Rngs(42), rngs=nnx.Rngs(42),
): ):
self.num_heads = num_heads self.num_heads = num_heads
self.in_features = in_features self.in_features = in_features
@@ -162,15 +159,15 @@ class MultiHeadAttention(nnx.Module):
if self.qkv_features % self.num_heads != 0: if self.qkv_features % self.num_heads != 0:
raise ValueError( raise ValueError(
f"Memory dimension ({self.qkv_features}) must be divisible by " f"Memory dimension ({self.qkv_features}) must be divisible by "
f"'num_heads' heads ({self.num_heads})." f"'num_heads' heads ({self.num_heads})."
) )
self.head_dim = self.qkv_features // self.num_heads self.head_dim = self.qkv_features // self.num_heads
linear_general = functools.partial( linear_general = functools.partial(
LinearGeneral, LinearGeneral,
out_features=(self.num_heads, self.head_dim), out_features=(self.num_heads, self.head_dim),
use_bias=self.use_bias, use_bias=self.use_bias,
) )
# project inputs_q to multi-headed q/k/v # project inputs_q to multi-headed q/k/v
# dimensions are then [batch..., length, n_heads, n_features_per_head] # dimensions are then [batch..., length, n_heads, n_features_per_head]
@@ -186,18 +183,18 @@ class MultiHeadAttention(nnx.Module):
self.key_ln = None self.key_ln = None
self.out = LinearGeneral( self.out = LinearGeneral(
in_features=(self.num_heads, self.head_dim), in_features=(self.num_heads, self.head_dim),
out_features=self.out_features, out_features=self.out_features,
axis=(-2, -1), axis=(-2, -1),
use_bias=self.use_bias, use_bias=self.use_bias,
rngs=rngs, rngs=rngs,
) )
self.use_per_dim_scale = use_per_dim_scale self.use_per_dim_scale = use_per_dim_scale
self.use_rotary_position_embeddings = use_rotary_position_embeddings self.use_rotary_position_embeddings = use_rotary_position_embeddings
if self.use_rotary_position_embeddings: if self.use_rotary_position_embeddings:
self.rotary_position_embedding = RotaryPositionalEmbedding( self.rotary_position_embedding = RotaryPositionalEmbedding(
embedding_dims=self.head_dim, embedding_dims=self.head_dim,
) )
else: else:
self.rotary_position_embedding = None self.rotary_position_embedding = None
@@ -208,20 +205,20 @@ class MultiHeadAttention(nnx.Module):
self.per_dim_scale = None self.per_dim_scale = None
def __call__( def __call__(
self, self,
inputs_q: Array, inputs_q: Array,
*, *,
decode_cache: DecodeCache | None = None, decode_cache: DecodeCache | None = None,
patch_mask: Array | None = None, patch_mask: Array | None = None,
deterministic: bool | None = None, deterministic: bool | None = None,
sow_weights: bool = False, sow_weights: bool = False,
) -> tuple[Float[Array, "b ... o"], DecodeCache | None]: ) -> tuple[Float[Array, "b ... o"], DecodeCache | None]:
"""Applies multi-head dot product attention on the input data.""" """Applies multi-head dot product attention on the input data."""
_, n_patches, input_in_features = inputs_q.shape _, n_patches, input_in_features = inputs_q.shape
if input_in_features != self.in_features: if input_in_features != self.in_features:
raise ValueError( raise ValueError(
f"Incompatible input dimension, got {input_in_features} " f"Incompatible input dimension, got {input_in_features} "
f"but module expects {self.in_features}." f"but module expects {self.in_features}."
) )
if patch_mask is None: if patch_mask is None:
patch_mask = jnp.zeros_like(inputs_q.shape[:-1], dtype=jnp.bool) 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) value = self.value(inputs_q)
if decode_cache is None: if decode_cache is None:
num_masked = jnp.sum( num_masked = jnp.sum(patch_mask.astype(jnp.int32), axis=-1, keepdims=False)
patch_mask.astype(jnp.int32), axis=-1, keepdims=False
)
next_index = jnp.zeros_like(num_masked, dtype=jnp.int32) next_index = jnp.zeros_like(num_masked, dtype=jnp.int32)
else: else:
num_masked = ( num_masked = (
jnp.sum(patch_mask.astype(jnp.int32), axis=-1, keepdims=False) jnp.sum(patch_mask.astype(jnp.int32), axis=-1, keepdims=False)
+ decode_cache.num_masked + decode_cache.num_masked
) )
next_index = decode_cache.next_index next_index = decode_cache.next_index
if self.use_rotary_position_embeddings: if self.use_rotary_position_embeddings:
position = ( position = (
jnp.arange(n_patches, dtype=jnp.int32)[None, :] jnp.arange(n_patches, dtype=jnp.int32)[None, :]
+ next_index[:, None] + next_index[:, None]
- num_masked[:, None] - num_masked[:, None]
) )
query = self.rotary_position_embedding(query, position) query = self.rotary_position_embedding(query, position)
key = self.rotary_position_embedding(key, 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.next_index = next_index + n_patches
decode_cache.num_masked = num_masked decode_cache.num_masked = num_masked
attn_mask = make_attn_mask( attn_mask = make_attn_mask(
query_length=n_patches, query_length=n_patches,
num_all_masked_kv=num_masked, num_all_masked_kv=num_masked,
query_index_offset=next_index, query_index_offset=next_index,
kv_length=decode_cache_size, kv_length=decode_cache_size,
) )
else: else:
# Training # Training
attn_mask = make_attn_mask( attn_mask = make_attn_mask(query_length=n_patches, num_all_masked_kv=num_masked)
query_length=n_patches, num_all_masked_kv=num_masked
)
# apply attention # apply attention
x = self.attention_fn( x = self.attention_fn(
query * jnp.sqrt(self.head_dim), query * jnp.sqrt(self.head_dim),
key, key,
value, value,
mask=attn_mask, mask=attn_mask,
deterministic=deterministic, deterministic=deterministic,
module=self if sow_weights else None, module=self if sow_weights else None,
) )
# back to the original inputs dimensions # back to the original inputs dimensions
out = self.out(x) out = self.out(x)
@@ -308,12 +301,12 @@ class Transformer(nnx.Module):
raise ValueError(f"Layer norm: {config.attention_norm} not supported.") raise ValueError(f"Layer norm: {config.attention_norm} not supported.")
self.attn = MultiHeadAttention( self.attn = MultiHeadAttention(
num_heads=config.num_heads, num_heads=config.num_heads,
in_features=config.model_dims, in_features=config.model_dims,
use_per_dim_scale=True, use_per_dim_scale=True,
use_rotary_position_embeddings=config.use_rotary_position_embeddings, use_rotary_position_embeddings=config.use_rotary_position_embeddings,
qk_norm=config.qk_norm, qk_norm=config.qk_norm,
rngs=rngs, rngs=rngs,
) )
if config.feedforward_norm == "rms": if config.feedforward_norm == "rms":
@@ -322,16 +315,16 @@ class Transformer(nnx.Module):
else: else:
raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.") raise ValueError(f"Layer norm: {config.feedforward_norm} not supported.")
self.ff0 = nnx.Linear( self.ff0 = nnx.Linear(
in_features=config.model_dims, in_features=config.model_dims,
out_features=config.hidden_dims, out_features=config.hidden_dims,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
self.ff1 = nnx.Linear( self.ff1 = nnx.Linear(
in_features=config.hidden_dims, in_features=config.hidden_dims,
out_features=config.model_dims, out_features=config.model_dims,
use_bias=config.use_bias, use_bias=config.use_bias,
rngs=rngs, rngs=rngs,
) )
if config.ff_activation == "relu": if config.ff_activation == "relu":
self.activation = jax.nn.relu self.activation = jax.nn.relu
@@ -343,23 +336,21 @@ class Transformer(nnx.Module):
raise ValueError(f"Activation: {config.ff_activation} not supported.") raise ValueError(f"Activation: {config.ff_activation} not supported.")
def __call__( def __call__(
self, self,
input_embeddings: Float[Array, "b n d"], input_embeddings: Float[Array, "b n d"],
patch_mask: Bool[Array, "b n"], patch_mask: Bool[Array, "b n"],
decode_cache: DecodeCache | None = None, decode_cache: DecodeCache | None = None,
) -> tuple[Float[Array, "b n d"], DecodeCache | None]: ) -> tuple[Float[Array, "b n d"], DecodeCache | None]:
attn_output, decode_cache = self.attn( attn_output, decode_cache = self.attn(
inputs_q=self.pre_attn_ln(input_embeddings), inputs_q=self.pre_attn_ln(input_embeddings),
decode_cache=decode_cache, decode_cache=decode_cache,
patch_mask=patch_mask, patch_mask=patch_mask,
sow_weights=False, sow_weights=False,
deterministic=True, deterministic=True,
) )
attn_output = self.post_attn_ln(attn_output) + input_embeddings attn_output = self.post_attn_ln(attn_output) + input_embeddings
output_embeddings = ( output_embeddings = (
self.post_ff_ln( self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))) + attn_output
)
+ attn_output
) )
return output_embeddings, decode_cache return output_embeddings, decode_cache
src/timesfm/flax/util.py
+25 -26
View File
@@ -19,7 +19,6 @@ import functools
import jax import jax
import jax.numpy as jnp import jax.numpy as jnp
import jaxtyping import jaxtyping
import typeguard
Float = jaxtyping.Float Float = jaxtyping.Float
Array = jaxtyping.Array Array = jaxtyping.Array
@@ -42,38 +41,38 @@ class DecodeCache:
@jax.jit @jax.jit
def update_running_stats( def update_running_stats(
n: Float[Array, "b"], n: Float[Array, "b"],
mu: Float[Array, "b"], mu: Float[Array, "b"],
sigma: Float[Array, "b"], sigma: Float[Array, "b"],
x: Float[Array, "b p"], x: Float[Array, "b p"],
mask: Bool[Array, "b p"], mask: Bool[Array, "b p"],
) -> tuple[ ) -> 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.""" """Updates the running stats."""
is_legit = jnp.logical_not(mask) is_legit = jnp.logical_not(mask)
inc_n = jnp.sum(is_legit.astype(jnp.float32), axis=-1, keepdims=False) inc_n = jnp.sum(is_legit.astype(jnp.float32), axis=-1, keepdims=False)
inc_mu = jnp.where( 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_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_n = n + inc_n
new_mu = jnp.where(new_n == 0, 0.0, (n * mu + inc_mu * inc_n) / new_n) new_mu = jnp.where(new_n == 0, 0.0, (n * mu + inc_mu * inc_n) / new_n)
new_sigma = jnp.sqrt( new_sigma = jnp.sqrt(
jnp.where( jnp.where(
new_n == 0, new_n == 0,
0.0, 0.0,
( (
n * sigma * sigma n * sigma * sigma
+ inc_n * inc_sigma * inc_sigma + inc_n * inc_sigma * inc_sigma
+ n * (mu - new_mu) * (mu - new_mu) + n * (mu - new_mu) * (mu - new_mu)
+ inc_n * (inc_mu - new_mu) * (inc_mu - new_mu) + inc_n * (inc_mu - new_mu) * (inc_mu - new_mu)
)
/ new_n,
) )
/ new_n,
)
) )
return (w := (new_n, new_mu, new_sigma), w) 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) 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) carry, moved_ys = jax.lax.scan(f, init, moved_xs, **kwargs)
return ( return (
carry, carry,
jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), moved_ys), jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), moved_ys),
) )
@functools.partial(jax.jit, static_argnames=("reverse",)) @functools.partial(jax.jit, static_argnames=("reverse",))
def revin( def revin(
x: Float[Array, "b ..."], x: Float[Array, "b ..."],
mu: Float[Array, "b ..."], mu: Float[Array, "b ..."],
sigma: Float[Array, "b ..."], sigma: Float[Array, "b ..."],
reverse: bool = False, reverse: bool = False,
): ):
"""Reversible per-instance normalization.""" """Reversible per-instance normalization."""
if len(mu.shape) == len(x.shape) - 1: if len(mu.shape) == len(x.shape) - 1:
src/timesfm/timesfm_2p5/timesfm_2p5_flax.py
+212 -216
View File
@@ -19,7 +19,9 @@ import functools
import gc import gc
import logging import logging
import math import math
from typing import Any, Callable import os
from pathlib import Path
from typing import Any, Callable, Dict
import einshape import einshape
from flax import nnx from flax import nnx
@@ -55,7 +57,7 @@ def try_gc():
@nnx.vmap(in_axes=(None, 0), out_axes=0) @nnx.vmap(in_axes=(None, 0), out_axes=0)
def _create_stacked_transformers( 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)) 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) 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) carry, moved_ys = jax.lax.scan(f, init, moved_xs, **kwargs)
return ( return (
carry, carry,
jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), moved_ys), 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)) @nnx.scan(in_axes=(0, nnx.Carry, None, 0), out_axes=(nnx.Carry, 0))
def _apply_stacked_transformers( def _apply_stacked_transformers(
model: transformer.Transformer, model: transformer.Transformer,
x: Float[Array, "b n d"], x: Float[Array, "b n d"],
m: Float[Array, "b n"], m: Float[Array, "b n"],
decode_cache: util.DecodeCache | None = None, decode_cache: util.DecodeCache | None = None,
) -> Float[Array, "b n d"]: ) -> Float[Array, "b n d"]:
return model(x, m, decode_cache=decode_cache) 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. # Layers.
self.tokenizer = dense.ResidualBlock(self.config.tokenizer) self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
self.stacked_xf = _create_stacked_transformers( self.stacked_xf = _create_stacked_transformers(
self.config.stacked_transformers, self.config.stacked_transformers,
jax.random.split(jax.random.key(42), self.x), jax.random.split(jax.random.key(42), self.x),
) )
self.output_projection_point = dense.ResidualBlock( self.output_projection_point = dense.ResidualBlock(
self.config.output_projection_point self.config.output_projection_point
) )
self.output_projection_quantiles = dense.ResidualBlock( self.output_projection_quantiles = dense.ResidualBlock(
self.config.output_projection_quantiles self.config.output_projection_quantiles
) )
def __call__( def __call__(
self, self,
inputs: Float[Array, "b n p"], inputs: Float[Array, "b n p"],
masks: Bool[Array, "b n p"], masks: Bool[Array, "b n p"],
decode_cache: util.DecodeCache | None = None, decode_cache: util.DecodeCache | None = None,
): ):
tokenizer_inputs = jnp.concatenate( tokenizer_inputs = jnp.concatenate([inputs, masks.astype(inputs.dtype)], axis=-1)
[inputs, masks.astype(inputs.dtype)], axis=-1
)
input_embeddings = self.tokenizer(tokenizer_inputs) input_embeddings = self.tokenizer(tokenizer_inputs)
if decode_cache is None: if decode_cache is None:
decode_cache = [None] * self.x decode_cache = [None] * self.x
output_embeddings, decode_cache = _apply_stacked_transformers( 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_ts = self.output_projection_point(output_embeddings)
output_quantile_spread = self.output_projection_quantiles(output_embeddings) output_quantile_spread = self.output_projection_quantiles(output_embeddings)
return ( return (
input_embeddings, input_embeddings,
output_embeddings, output_embeddings,
output_ts, output_ts,
output_quantile_spread, output_quantile_spread,
), decode_cache ), decode_cache
@nnx.jit(static_argnames=("horizon",)) @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_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) 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( (last_n, last_mu, last_sigma), (_, context_mu, context_sigma) = scan(
lambda carry, xs: util.update_running_stats(*carry, *xs), lambda carry, xs: util.update_running_stats(*carry, *xs),
init=(zero := jnp.zeros(shape=(batch_size)), zero, zero), init=(zero := jnp.zeros(shape=(batch_size)), zero, zero),
xs=(patched_inputs, patched_masks), xs=(patched_inputs, patched_masks),
axis=1, axis=1,
) )
decode_cache = util.DecodeCache( decode_cache = util.DecodeCache(
next_index=jnp.zeros(shape=(self.x, batch_size), dtype=jnp.int32), next_index=jnp.zeros(shape=(self.x, batch_size), dtype=jnp.int32),
num_masked=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( key=jnp.zeros(shape=(self.x, batch_size, decode_cache_size, self.h, self.hd)),
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)),
),
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 = revin(patched_inputs, context_mu, context_sigma, reverse=False)
normed_inputs = jnp.where(patched_masks, 0.0, normed_inputs) normed_inputs = jnp.where(patched_masks, 0.0, normed_inputs)
(_, _, normed_outputs, normed_quantile_spread), decode_cache = self( (_, _, normed_outputs, normed_quantile_spread), decode_cache = self(
normed_inputs, patched_masks, decode_cache normed_inputs, patched_masks, decode_cache
) )
renormed_outputs = jax_einshape( renormed_outputs = jax_einshape(
"bn(oq)->bnoq", "bn(oq)->bnoq",
revin(normed_outputs, context_mu, context_sigma, reverse=True), revin(normed_outputs, context_mu, context_sigma, reverse=True),
o=self.o, o=self.o,
q=self.q, q=self.q,
) )
renormed_quantile_spread = jax_einshape( renormed_quantile_spread = jax_einshape(
"bn(oq)->bnoq", "bn(oq)->bnoq",
revin(normed_quantile_spread, context_mu, context_sigma, reverse=True), revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
o=self.os, o=self.os,
q=self.q, q=self.q,
)[:, -1, ...] )[:, -1, ...]
# Autogressive decode # 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): def _ar_decode(module, carry, unused_iter):
last_renormed_output, (last_n, last_mu, last_sigma), decode_cache = carry last_renormed_output, (last_n, last_mu, last_sigma), decode_cache = carry
new_patched_input = jax_einshape( 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) new_mask = jnp.zeros_like(new_patched_input, dtype=jnp.bool)
carry_stats, (_, new_mu, new_sigma) = scan( carry_stats, (_, new_mu, new_sigma) = scan(
lambda carry, xs: util.update_running_stats(*carry, *xs), lambda carry, xs: util.update_running_stats(*carry, *xs),
init=(last_n, last_mu, last_sigma), init=(last_n, last_mu, last_sigma),
xs=(new_patched_input, new_mask), xs=(new_patched_input, new_mask),
axis=1, axis=1,
)
new_normed_input = revin(
new_patched_input, new_mu, new_sigma, reverse=False
) )
new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False)
(_, _, new_normed_output, _), decode_cache = module( (_, _, 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( new_renormed_output = jax_einshape(
"bm(oq)->bmoq", "bm(oq)->bmoq",
revin(new_normed_output, new_mu, new_sigma, reverse=True), revin(new_normed_output, new_mu, new_sigma, reverse=True),
o=module.o, o=module.o,
q=module.q, q=module.q,
)[..., -1, :, :] )[..., -1, :, :]
return ( return (
( (
new_renormed_output[..., module.decode_index], new_renormed_output[..., module.decode_index],
carry_stats, carry_stats,
decode_cache, decode_cache,
), ),
new_renormed_output, new_renormed_output,
) )
if num_decode_steps > 0: if num_decode_steps > 0:
_, ar_renormed_outputs = _ar_decode( _, ar_renormed_outputs = _ar_decode(
self, self,
( (
renormed_outputs[..., -1, :, self.decode_index], renormed_outputs[..., -1, :, self.decode_index],
(last_n, last_mu, last_sigma), (last_n, last_mu, last_sigma),
decode_cache, decode_cache,
), ),
jnp.arange(num_decode_steps), jnp.arange(num_decode_steps),
) )
else: else:
ar_renormed_outputs = None 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 return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs
def compile( def compile(
self, self,
context: int, context: int,
horizon: int, horizon: int,
per_core_batch_size: int = 1, per_core_batch_size: int = 1,
): ):
if context % self.p != 0: if context % self.p != 0:
logging.info( logging.info(
"When compiling, context needs to be multiple of the patch size %d." "When compiling, context needs to be multiple of the patch size %d."
" Modifying context to %d.", " Modifying context to %d.",
self.p, self.p,
context := math.ceil(context / self.p) * self.p, context := math.ceil(context / self.p) * self.p,
) )
if horizon % self.o != 0: if horizon % self.o != 0:
logging.info( logging.info(
"When compiling, horizon needs to be multiple of the output patch" "When compiling, horizon needs to be multiple of the output patch"
" size %d. Modifying horizon to %d.", " size %d. Modifying horizon to %d.",
self.o, self.o,
horizon := math.ceil(horizon / self.o) * self.o, horizon := math.ceil(horizon / self.o) * self.o,
) )
self.context = context 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 self.per_core_batch_size = per_core_batch_size
@nnx.pmap( @nnx.pmap(
in_axes=(None, None, 0, 0), in_axes=(None, None, 0, 0),
out_axes=(0, 0, 0), out_axes=(0, 0, 0),
devices=jax.devices(self.backend), devices=jax.devices(self.backend),
axis_size=self.num_devices, axis_size=self.num_devices,
static_broadcasted_argnums=(1,), static_broadcasted_argnums=(1,),
axis_name="global_batch", axis_name="global_batch",
) )
def compiled_decode_kernel(model, horizon, inputs, masks): def compiled_decode_kernel(model, horizon, inputs, masks):
return model.decode(horizon, inputs, masks) return model.decode(horizon, inputs, masks)
@@ -286,27 +278,23 @@ def _flip_quantile_fn(x):
@functools.partial( @functools.partial(
jax.jit, jax.jit,
donate_argnums=(0, 1, 2), donate_argnums=(0, 1, 2),
) )
def _force_flip_invariance_fn( def _force_flip_invariance_fn(
flipped_pf_outputs, flipped_pf_outputs,
flipped_quantile_spreads, flipped_quantile_spreads,
flipped_ar_outputs, flipped_ar_outputs,
): ):
"""Forces flip invariance.""" """Forces flip invariance."""
flipped_pf_outputs = _flip_quantile_fn(flipped_pf_outputs) flipped_pf_outputs = _flip_quantile_fn(flipped_pf_outputs)
flipped_pf_outputs = jax_einshape("tb...->(tb)...", 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 = _flip_quantile_fn(flipped_quantile_spreads)
flipped_quantile_spreads = jax_einshape( flipped_quantile_spreads = jax_einshape("tb...->(tb)...", flipped_quantile_spreads)
"tb...->(tb)...", flipped_quantile_spreads
)
to_concat = [flipped_pf_outputs[:, -1, ...]] to_concat = [flipped_pf_outputs[:, -1, ...]]
if flipped_ar_outputs is not None: if flipped_ar_outputs is not None:
flipped_ar_outputs = _flip_quantile_fn(flipped_ar_outputs) flipped_ar_outputs = _flip_quantile_fn(flipped_ar_outputs)
flipped_ar_outputs = jax_einshape( flipped_ar_outputs = jax_einshape("tbno...->(tb)(no)...", flipped_ar_outputs)
"tbno...->(tb)(no)...", flipped_ar_outputs
)
to_concat.append(flipped_ar_outputs) to_concat.append(flipped_ar_outputs)
flipped_full_forecast = jnp.concatenate(to_concat, axis=1) flipped_full_forecast = jnp.concatenate(to_concat, axis=1)
@@ -314,27 +302,25 @@ def _force_flip_invariance_fn(
@functools.partial( @functools.partial(
jax.jit, jax.jit,
static_argnames=("max_horizon",), static_argnames=("max_horizon",),
donate_argnums=(0,), donate_argnums=(0,),
) )
def _use_continuous_quantile_head_fn( def _use_continuous_quantile_head_fn(full_forecast, quantile_spreads, max_horizon):
full_forecast, quantile_spreads, max_horizon
):
"""Uses continuous quantile head.""" """Uses continuous quantile head."""
to_stack = [full_forecast[..., :max_horizon, 0]] to_stack = [full_forecast[..., :max_horizon, 0]]
for quantile_index in [1, 2, 3, 4]: for quantile_index in [1, 2, 3, 4]:
to_stack.append( to_stack.append(
quantile_spreads[:, :max_horizon, quantile_index] quantile_spreads[:, :max_horizon, quantile_index]
- quantile_spreads[:, :max_horizon, 5] - quantile_spreads[:, :max_horizon, 5]
+ full_forecast[:, :max_horizon, 5] + full_forecast[:, :max_horizon, 5]
) )
to_stack.append(full_forecast[..., :max_horizon, 5]) to_stack.append(full_forecast[..., :max_horizon, 5])
for quantile_index in [6, 7, 8, 9]: for quantile_index in [6, 7, 8, 9]:
to_stack.append( to_stack.append(
quantile_spreads[:, :max_horizon, quantile_index] quantile_spreads[:, :max_horizon, quantile_index]
- quantile_spreads[:, :max_horizon, 5] - quantile_spreads[:, :max_horizon, 5]
+ full_forecast[:, :max_horizon, 5] + full_forecast[:, :max_horizon, 5]
) )
return jnp.stack(to_stack, axis=-1) return jnp.stack(to_stack, axis=-1)
@@ -343,27 +329,27 @@ def _use_continuous_quantile_head_fn(
def _fix_quantile_crossing_fn(full_forecast): def _fix_quantile_crossing_fn(full_forecast):
"""Fixes quantile crossing.""" """Fixes quantile crossing."""
lower_quantiles = _scan_along_axis( lower_quantiles = _scan_along_axis(
lambda carry, x: (w := jnp.minimum(carry, x), w), lambda carry, x: (w := jnp.minimum(carry, x), w),
init=full_forecast[..., 5], init=full_forecast[..., 5],
xs=full_forecast[..., 1:5], xs=full_forecast[..., 1:5],
axis=-1, axis=-1,
reverse=True, reverse=True,
)[1] )[1]
upper_quantiles = _scan_along_axis( upper_quantiles = _scan_along_axis(
lambda carry, x: (w := jnp.maximum(carry, x), w), lambda carry, x: (w := jnp.maximum(carry, x), w),
init=full_forecast[..., 5], init=full_forecast[..., 5],
xs=full_forecast[..., 6:10], xs=full_forecast[..., 6:10],
axis=-1, axis=-1,
reverse=False, reverse=False,
)[1] )[1]
return jnp.concatenate( return jnp.concatenate(
[ [
full_forecast[..., :1], full_forecast[..., :1],
lower_quantiles, lower_quantiles,
full_forecast[..., 5:6], full_forecast[..., 5:6],
upper_quantiles, upper_quantiles,
], ],
axis=-1, axis=-1,
) )
@@ -389,25 +375,25 @@ def _before_model_decode(fc, inputs, masks):
@functools.partial( @functools.partial(
jax.jit, jax.jit,
static_argnames=( static_argnames=(
"fc", "fc",
"p", "p",
), ),
donate_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9), donate_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9),
) )
def _after_model_decode( def _after_model_decode(
fc, fc,
pf_outputs, pf_outputs,
quantile_spreads, quantile_spreads,
ar_outputs, ar_outputs,
flipped_pf_outputs, flipped_pf_outputs,
flipped_quantile_spreads, flipped_quantile_spreads,
flipped_ar_outputs, flipped_ar_outputs,
is_positive, is_positive,
mu, mu,
sigma, sigma,
p, p,
): ):
"""All Jax steps after model decode call.""" """All Jax steps after model decode call."""
# t: num_devices, b: per_core_batch_size # t: num_devices, b: per_core_batch_size
@@ -421,11 +407,11 @@ def _after_model_decode(
if fc.force_flip_invariance: if fc.force_flip_invariance:
( (
flipped_quantile_spreads, flipped_quantile_spreads,
flipped_pf_outputs, flipped_pf_outputs,
flipped_full_forecast, flipped_full_forecast,
) = _force_flip_invariance_fn( ) = _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 quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2
pf_outputs = (pf_outputs - flipped_pf_outputs) / 2 pf_outputs = (pf_outputs - flipped_pf_outputs) / 2
@@ -433,7 +419,7 @@ def _after_model_decode(
if fc.use_continuous_quantile_head: if fc.use_continuous_quantile_head:
full_forecast = _use_continuous_quantile_head_fn( 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: if fc.return_backcast:
@@ -448,9 +434,9 @@ def _after_model_decode(
if is_positive is not None: if is_positive is not None:
full_forecast = jnp.where( full_forecast = jnp.where(
is_positive[..., None], is_positive[..., None],
jnp.maximum(full_forecast, jnp.zeros_like(full_forecast)), jnp.maximum(full_forecast, jnp.zeros_like(full_forecast)),
full_forecast, full_forecast,
) )
return full_forecast return full_forecast
@@ -463,117 +449,127 @@ class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
@classmethod @classmethod
def from_pretrained( def from_pretrained(
cls, cls,
*, model_id: str = "google/timesfm-2.5-200m-flax",
path: str | None = None, *,
hf_repo_id: str | None = "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.""" """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)
else:
raise ValueError("Either path or hf_repo_id must be provided.")
instance = cls() # 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() checkpointer = ocp.StandardCheckpointer()
graph, state = nnx.split(instance.model) graph, state = nnx.split(instance.model)
state = checkpointer.restore(path, state) state = checkpointer.restore(model_file_path, state)
instance.model = nnx.merge(graph, state) instance.model = nnx.merge(graph, state)
return instance return instance
def compile(self, forecast_config: configs.ForecastConfig, **kwargs): def compile(self, forecast_config: configs.ForecastConfig, **kwargs):
# Acrobym used during validation. # Acrobym used during validation.
fc = forecast_config fc = forecast_config
if fc.max_context % self.model.p != 0: if fc.max_context % self.model.p != 0:
logging.info( logging.info(
"When compiling, max context needs to be multiple of the patch size" "When compiling, max context needs to be multiple of the patch size"
" %d. Using max context = %d instead.", " %d. Using max context = %d instead.",
self.model.p, self.model.p,
new_context := math.ceil(fc.max_context / self.model.p) new_context := math.ceil(fc.max_context / self.model.p) * self.model.p,
* self.model.p,
) )
fc = dataclasses.replace(fc, max_context=new_context) fc = dataclasses.replace(fc, max_context=new_context)
if fc.max_horizon % self.model.o != 0: if fc.max_horizon % self.model.o != 0:
logging.info( logging.info(
"When compiling, max horizon needs to be multiple of the output patch" "When compiling, max horizon needs to be multiple of the output patch"
" size %d. Using max horizon = %d instead.", " size %d. Using max horizon = %d instead.",
self.model.o, self.model.o,
new_horizon := math.ceil(fc.max_horizon / self.model.o) new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o,
* self.model.o,
) )
forecast_config = dataclasses.replace(fc, max_horizon=new_horizon) forecast_config = dataclasses.replace(fc, max_horizon=new_horizon)
if fc.max_context + fc.max_horizon > self.model.config.context_limit: if fc.max_context + fc.max_horizon > self.model.config.context_limit:
raise ValueError( raise ValueError(
"Context + horizon must be less than the context limit." "Context + horizon must be less than the context limit."
f" {fc.max_context} + {fc.max_horizon} >" f" {fc.max_context} + {fc.max_horizon} >"
f" {self.model.config.context_limit}." f" {self.model.config.context_limit}."
) )
if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os): if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os):
raise ValueError( raise ValueError(
"Continuous quantile head is not supported for horizons >" f"Continuous quantile head is not supported for horizons > {self.model.os}."
f" {self.model.os}."
) )
self.forecast_config = forecast_config self.forecast_config = forecast_config
self.model.compile( self.model.compile(
context=self.forecast_config.max_context, context=self.forecast_config.max_context,
horizon=self.forecast_config.max_horizon, horizon=self.forecast_config.max_horizon,
per_core_batch_size=fc.per_core_batch_size, per_core_batch_size=fc.per_core_batch_size,
) )
self.per_core_batch_size = self.forecast_config.per_core_batch_size self.per_core_batch_size = self.forecast_config.per_core_batch_size
self.num_devices = self.model.num_devices self.num_devices = self.model.num_devices
self.global_batch_size = ( 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): def compiled_decode_kernel(fc, horizon, inputs, masks):
inputs = jnp.array(inputs, dtype=jnp.float32) inputs = jnp.array(inputs, dtype=jnp.float32)
masks = jnp.array(masks, dtype=jnp.bool) masks = jnp.array(masks, dtype=jnp.bool)
if horizon > fc.max_horizon: if horizon > fc.max_horizon:
raise ValueError( raise ValueError(
"Horizon must be less than the max horizon." f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}."
f" {horizon} > {fc.max_horizon}."
) )
to_trim = fc.max_horizon - horizon to_trim = fc.max_horizon - horizon
inputs, masks, is_positive, mu, sigma = _before_model_decode( inputs, masks, is_positive, mu, sigma = _before_model_decode(fc, inputs, masks)
fc, inputs, masks
)
pf_outputs, quantile_spreads, ar_outputs = self.model.compiled_decode( 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: if fc.force_flip_invariance:
flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( 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: else:
flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = ( flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
None, None,
None, None,
None, None,
) )
full_forecast = _after_model_decode( full_forecast = _after_model_decode(
fc, fc,
pf_outputs, pf_outputs,
quantile_spreads, quantile_spreads,
ar_outputs, ar_outputs,
flipped_pf_outputs, flipped_pf_outputs,
flipped_quantile_spreads, flipped_quantile_spreads,
flipped_ar_outputs, flipped_ar_outputs,
is_positive, is_positive,
mu, mu,
sigma, sigma,
self.model.p, self.model.p,
) )
full_forecast_np = np.array(full_forecast) full_forecast_np = np.array(full_forecast)
del 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 return full_forecast_np[..., 5], full_forecast_np
self.compiled_decode = functools.partial( 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. # Layers.
self.tokenizer = dense.ResidualBlock(self.config.tokenizer) self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
self.stacked_xf = nn.ModuleList( self.stacked_xf = nn.ModuleList([
[ transformer.Transformer(self.config.stacked_transformers.transformer)
transformer.Transformer(self.config.stacked_transformers.transformer) for _ in range(self.x)
for _ in range(self.x) ])
]
)
self.output_projection_point = dense.ResidualBlock( self.output_projection_point = dense.ResidualBlock(
self.config.output_projection_point self.config.output_projection_point
) )
@@ -272,7 +270,7 @@ class TimesFM_2p5_200M_torch(timesfm_2p5_base.TimesFM_2p5, ModelHubMixin):
def _from_pretrained( def _from_pretrained(
cls, cls,
*, *,
model_id: str, model_id: str = "google/timesfm-2.5-200m-pytorch",
revision: Optional[str], revision: Optional[str],
cache_dir: Optional[Union[str, Path]], cache_dir: Optional[Union[str, Path]],
force_download: bool, force_download: bool,
where
View File