FEDformerConfig

model.FEDformerConfig

Hyperparameters for :class:FEDformerModel.

__init__(self, h=24, input_size=96, hidden_size=128, n_heads=8, modes=64, mode_select='random', moving_avg_window=25, encoder_layers=2, decoder_layers=1, conv_hidden_size=32, decoder_input_size_multiplier=0.5, dropout=0.05, activation='gelu', fea_activation='tanh', random_seed=1)

Parameter	Type	Default	Description
h	int	24	Forecast horizon (number of future steps to predict).
input_size	int	96	Length of the historical context fed to the encoder.
hidden_size	int	128	Embedding / attention hidden dimension.
n_heads	int	8	Number of heads. FEDformer requires exactly 8 (its spectral weight tensor is built with a leading dimension of 8).
modes	int	64	Number of Fourier modes kept by each spectral block.
mode_select	str	"random"	Mode-selection strategy: "random" or "lowest".
moving_avg_window	int	25	Kernel size for the trend moving-average filter.
encoder_layers	int	2	Number of stacked encoder layers.
decoder_layers	int	1	Number of stacked decoder layers.
conv_hidden_size	int	32	Hidden channels of the position-wise conv FFN.
decoder_input_size_multiplier	float	0.5	Fraction of input_size used as the decoder start-token ("label") length; must be in (0, 1).
dropout	float	0.05	Dropout rate applied throughout (training only).
activation	str	"gelu"	Conv-FFN nonlinearity -- "relu" or "gelu".
fea_activation	str	"tanh"	Frequency cross-attention score nonlinearity -- "tanh" or "softmax".
random_seed	int	1	Base seed for deterministic Fourier-mode selection.

FEDformerModel

model.FEDformerModel · inherits flax.linen.Module

Univariate FEDformer forecaster built with flax.linen.

Forward pass: [B, input_size, 1] -> [B, h, 1].

During training pass deterministic=False and supply a "dropout" RNG via rngs={"dropout": key} in model.apply(...).

__init__(self, config)

Parameter	Type	Default	Description
config	FEDformerConfig	-	(undocumented)

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

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

Returns: jnp.ndarray (The forecast output, shape [B, h, 1]).

FourierBlock

model.FourierBlock · inherits flax.linen.Module

Frequency Enhanced Block (FEB-f) -- self-attention in Fourier space.

__init__(self, in_channels, out_channels, seq_len, modes=64, mode_select='random', mode_seed=0)

Parameter	Type	Default	Description
in_channels	int	-	Total input feature width (== hidden_size).
out_channels	int	-	Total output feature width (== hidden_size).
seq_len	int	-	Sequence length used to pick the frequency modes.
modes	int	64	Number of Fourier modes to keep.
mode_select	str	"random"	"random" or "lowest".
mode_seed	int	0	Seed for deterministic random mode selection.

__call__(self, q, k, v, deterministic=True) -> jnp.ndarray

Parameter	Type	Default	Description
q	jnp.ndarray	-	(undocumented)
k	jnp.ndarray	-	unused (self-attention reads only the query stream)
v	jnp.ndarray	-	unused
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).

FourierCrossAttention

model.FourierCrossAttention · inherits flax.linen.Module

Frequency Enhanced Attention (FEA-f) -- cross-attention in Fourier space.

__init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes=64, mode_select='random', activation='tanh', mode_seed=0)

Parameter	Type	Default	Description
in_channels	int	-	Input feature width (== hidden_size).
out_channels	int	-	Output feature width (== hidden_size).
seq_len_q	int	-	Decoder (query) sequence length, for query-mode selection.
seq_len_kv	int	-	Encoder (key/value) sequence length, for key-mode selection.
modes	int	64	Number of Fourier modes to keep on each side.
mode_select	str	"random"	"random" or "lowest".
activation	str	"tanh"	Frequency-domain score nonlinearity: "tanh" or "softmax".
mode_seed	int	0	Seed for deterministic random mode selection.

__call__(self, q, k, v, deterministic=True) -> jnp.ndarray

Parameter	Type	Default	Description
q	jnp.ndarray	-	decoder stream [B, Lq, H, E]
k	jnp.ndarray	-	encoder stream [B, Lk, H, E]
v	jnp.ndarray	-	encoder stream (unused; FEA reuses k as values)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).

MultiHeadProjection

model.MultiHeadProjection · inherits flax.linen.Module

Wrap a spectral block with multi-head Q/K/V/output projections.

__init__(self, inner, hidden_size, n_heads)

Parameter	Type	Default	Description
inner	fnn.Module	-	The spectral mixing module (FourierBlock or FourierCrossAttention). Its __call__ must accept (q, k, v, deterministic) with the per-head layout [B, L, H, E].
hidden_size	int	-	Total feature width.
n_heads	int	-	Number of heads (must be 8 for FEDformer).

__call__(self, queries, keys, values, deterministic=True) -> jnp.ndarray

Parameter	Type	Default	Description
queries	jnp.ndarray	-	(undocumented)
keys	jnp.ndarray	-	(undocumented)
values	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (undocumented).

SeasonalLayerNorm

model.SeasonalLayerNorm · inherits flax.linen.Module

LayerNorm then remove the time-mean (keeps the seasonal part zero-mean).

__init__(self, hidden_size)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)

__call__(self, x) -> jnp.ndarray

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

Returns: jnp.ndarray (undocumented).

TokenEmbedding

model.TokenEmbedding · inherits flax.linen.Module

Conv1D(k=3, circular) value embedding + dropout.

__init__(self, hidden_size, dropout_rate)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
dropout_rate	float	-	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

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

Returns: jnp.ndarray (undocumented).

EncoderLayer

model.EncoderLayer · inherits flax.linen.Module

FEDformer encoder layer (FEB self-attention + conv FFN, both decomposed).

__init__(self, hidden_size, conv_hidden_size, n_heads, moving_avg_window, dropout_rate, activation, seq_len, modes, mode_select, mode_seed)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
conv_hidden_size	int	-	(undocumented)
n_heads	int	-	(undocumented)
moving_avg_window	int	-	(undocumented)
dropout_rate	float	-	(undocumented)
activation	str	-	(undocumented)
seq_len	int	-	(undocumented)
modes	int	-	(undocumented)
mode_select	str	-	(undocumented)
mode_seed	int	-	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

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

Returns: jnp.ndarray (undocumented).

Encoder

model.Encoder · inherits flax.linen.Module

Stack of EncoderLayers followed by a seasonal LayerNorm.

__init__(self, n_layers, hidden_size, conv_hidden_size, n_heads, moving_avg_window, dropout_rate, activation, seq_len, modes, mode_select, mode_seed)

Parameter	Type	Default	Description
n_layers	int	-	(undocumented)
hidden_size	int	-	(undocumented)
conv_hidden_size	int	-	(undocumented)
n_heads	int	-	(undocumented)
moving_avg_window	int	-	(undocumented)
dropout_rate	float	-	(undocumented)
activation	str	-	(undocumented)
seq_len	int	-	(undocumented)
modes	int	-	(undocumented)
mode_select	str	-	(undocumented)
mode_seed	int	-	(undocumented)

__call__(self, x, deterministic=True) -> jnp.ndarray

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

Returns: jnp.ndarray (undocumented).

DecoderLayer

model.DecoderLayer · inherits flax.linen.Module

FEDformer decoder layer (self-FEB, cross-FEA, conv FFN; trend extracted).

__init__(self, hidden_size, conv_hidden_size, n_heads, c_out, moving_avg_window, dropout_rate, activation, self_seq_len, cross_seq_len_q, cross_seq_len_kv, modes, mode_select, fea_activation, mode_seed)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
conv_hidden_size	int	-	(undocumented)
n_heads	int	-	(undocumented)
c_out	int	-	(undocumented)
moving_avg_window	int	-	(undocumented)
dropout_rate	float	-	(undocumented)
activation	str	-	(undocumented)
self_seq_len	int	-	(undocumented)
cross_seq_len_q	int	-	(undocumented)
cross_seq_len_kv	int	-	(undocumented)
modes	int	-	(undocumented)
mode_select	str	-	(undocumented)
fea_activation	str	-	(undocumented)
mode_seed	int	-	(undocumented)

__call__(self, x, cross, deterministic=True) -> Tuple[jnp.ndarray, jnp.ndarray]

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

Returns: Tuple[jnp.ndarray, jnp.ndarray] (undocumented).

Decoder

model.Decoder · inherits flax.linen.Module

Stack of DecoderLayers + seasonal LayerNorm + final linear projection.

__init__(self, n_layers, hidden_size, conv_hidden_size, n_heads, c_out, moving_avg_window, dropout_rate, activation, self_seq_len, cross_seq_len_q, cross_seq_len_kv, modes, mode_select, fea_activation, mode_seed)

Parameter	Type	Default	Description
n_layers	int	-	(undocumented)
hidden_size	int	-	(undocumented)
conv_hidden_size	int	-	(undocumented)
n_heads	int	-	(undocumented)
c_out	int	-	(undocumented)
moving_avg_window	int	-	(undocumented)
dropout_rate	float	-	(undocumented)
activation	str	-	(undocumented)
self_seq_len	int	-	(undocumented)
cross_seq_len_q	int	-	(undocumented)
cross_seq_len_kv	int	-	(undocumented)
modes	int	-	(undocumented)
mode_select	str	-	(undocumented)
fea_activation	str	-	(undocumented)
mode_seed	int	-	(undocumented)

__call__(self, x, cross, trend, deterministic=True) -> Tuple[jnp.ndarray, jnp.ndarray]

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

Returns: Tuple[jnp.ndarray, jnp.ndarray] (undocumented).