EnEmbedding

timexer_module.EnEmbedding · inherits nnx.Module

Endogenous patch embedding with a learnable per-variate global token.

x [B, N, L] is cut into L // patch_len non-overlapping patches (any tail shorter than a patch is dropped, as in the reference's unfold), each projected by a bias-free Linear and shifted by the sinusoidal positional code of its patch index; the global token is appended LAST. Returns ([B * N, patch_num + 1, d_model], n_vars).

__init__(self, n_vars: int, d_model: int, patch_len: int, dropout: float, *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
n_vars	int	-	(undocumented)
d_model	int	-	(undocumented)
patch_len	int	-	(undocumented)
dropout	float	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray, *, deterministic: bool) -> tuple[jnp.ndarray, int]

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: tuple[jnp.ndarray, int] (undocumented)

DataEmbeddingInverted

timexer_module.DataEmbeddingInverted · inherits nnx.Module

Variate-as-token embedding: each variate's full window is one token, projected time -> hidden. x [B, L, N] -> [B, N, d_model].

__init__(self, c_in: int, d_model: int, dropout: float, *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
c_in	int	-	(undocumented)
d_model	int	-	(undocumented)
dropout	float	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray, *, deterministic: bool) -> jnp.ndarray

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: jnp.ndarray (undocumented)

AttentionLayer

timexer_module.AttentionLayer · inherits nnx.Module

Multi-head full (softmax dot-product) attention, self- or cross-.

Folds the reference AttentionLayer + FullAttention (no mask, weights-dropout, scale d_k ** -0.5). q from q_in; k/v from kv_in (equal for self-attention). torch nn.Linear init.

__init__(self, *, hidden_size: int, n_heads: int, attn_dropout: float, rngs: nnx.Rngs)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
n_heads	int	-	(undocumented)
attn_dropout	float	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, q_in, kv_in, *, deterministic: bool)

Parameter	Type	Default	Description
q_in	-	-	(undocumented)
kv_in	-	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: (undocumented)

TimeXerEncoderLayer

timexer_module.TimeXerEncoderLayer · inherits nnx.Module

One TimeXer encoder layer: patch-token self-attention, global-token-only cross-attention into the variate context, position-wise FFN.

x [B * n_vars, pn + 1, d] with the global token LAST; cross is [B, n_cross_tokens, d]. Only the global token row reads the cross context inside a layer — patch tokens see it through later layers.

__init__(self, *, hidden_size: int, n_heads: int, d_ff: int, dropout: float, rngs: nnx.Rngs)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
n_heads	int	-	(undocumented)
d_ff	int	-	(undocumented)
dropout	float	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray, cross: jnp.ndarray, *, n_vars: int, deterministic: bool) -> jnp.ndarray

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
cross	jnp.ndarray	-	(undocumented)
n_vars	int	-	(undocumented)
deterministic	bool	-	(undocumented)

Returns: jnp.ndarray (undocumented)

TimeXerNet

timexer_module.TimeXerNet · inherits nnx.Module

Full TimeXer: NS-norm -> patch + variate embeddings -> encoder stack -> per-variate flatten head -> NS-denorm.

__call__(insample_y [B, L, N], deterministic) -> [B, h, N * mult] (the reference reshapes its [B, h * mult, N] head output with reshape(B, h, -1) — kept verbatim).

__init__(self, *, h: int, input_size: int, n_series: int, patch_len: int = 16, hidden_size: int = 512, n_heads: int = 8, e_layers: int = 2, d_ff: int = 2048, dropout: float = 0.1, use_norm: bool = True, outputsize_multiplier: int = 1, rngs: nnx.Rngs)

Parameter	Type	Default	Description
h	int	-	(undocumented)
input_size	int	-	(undocumented)
n_series	int	-	(undocumented)
patch_len	int	16	(undocumented)
hidden_size	int	512	(undocumented)
n_heads	int	8	(undocumented)
e_layers	int	2	(undocumented)
d_ff	int	2048	(undocumented)
dropout	float	0.1	(undocumented)
use_norm	bool	True	(undocumented)
outputsize_multiplier	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, insample_y: jnp.ndarray, *, deterministic: bool = True) -> jnp.ndarray

Parameter	Type	Default	Description
insample_y	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented)