The following reference documents the public API for the Chronax FEDformer implementation.
FEDformer
chronax.models.FEDformer · inherits BaseForecaster
Univariate FEDformer forecaster (JAX/Flax Linen port of neuralforecast.FEDformer).
Maintenance status: Active univariate forecaster. Integrates with the BaseForecaster interface, including conformal prediction intervals via predict(level=...). Defaults match neuralforecast.FEDformer architecture (n_heads must be 8); input_size=-1 resolves to 3 * h.
Attributes
uses_exog: False
alias: "FEDformer"
conformal_params: ConformalIntervals | None
model_: TrainState | None
__init__(self, h, input_size=-1, 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', max_steps=1000, learning_rate=0.0001, batch_size=32, num_lr_decays=3, val_fraction=0.1, val_check_steps=100, early_stop_patience_steps=-1, grad_clip=1.0, weight_decay=0.0, random_seed=1, alias='FEDformer', loss='mae')
Initializes the FEDformer model configuration and training parameters.
| Parameter | Type | Default | Description |
|---|---|---|---|
h |
int |
- | (undocumented) |
input_size |
int |
-1 |
(undocumented) |
hidden_size |
int |
128 |
(undocumented) |
n_heads |
int |
8 |
(undocumented) |
modes |
int |
64 |
(undocumented) |
mode_select |
str |
"random" |
(undocumented) |
moving_avg_window |
int |
25 |
(undocumented) |
encoder_layers |
int |
2 |
(undocumented) |
decoder_layers |
int |
1 |
(undocumented) |
conv_hidden_size |
int |
32 |
(undocumented) |
decoder_input_size_multiplier |
float |
0.5 |
(undocumented) |
dropout |
float |
0.05 |
(undocumented) |
activation |
str |
"gelu" |
(undocumented) |
fea_activation |
str |
"tanh" |
(undocumented) |
max_steps |
int |
1000 |
(undocumented) |
learning_rate |
float |
1e-4 |
(undocumented) |
batch_size |
int |
32 |
(undocumented) |
num_lr_decays |
int |
3 |
(undocumented) |
val_fraction |
float |
0.1 |
(undocumented) |
val_check_steps |
int |
100 |
(undocumented) |
early_stop_patience_steps |
int |
-1 |
(undocumented) |
grad_clip |
float |
1.0 |
(undocumented) |
weight_decay |
float |
0.0 |
(undocumented) |
random_seed |
int |
1 |
(undocumented) |
alias |
str |
"FEDformer" |
(undocumented) |
loss |
Union[str, LossFn] |
"mae" |
(undocumented) |
fit(self, y, X=None) -> Self
Fit on a univariate 1-D series.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
y |
jnp.ndarray |
- | (undocumented) |
X |
jnp.ndarray \| None |
None |
(undocumented) |
Returns: Self (the fitted forecaster; sets self.model_).
Raises:
* NotImplementedError: Exogenous variables are not supported.
* ValueError: If y is not 1-D or is too short for the configured input_size and h.
predict(self, h, X=None, level=None) -> dict
Forecast h steps from the fitted context.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
h |
int |
- | (undocumented) |
X |
jnp.ndarray \| None |
None |
(undocumented) |
level |
list[int \| float] \| None |
None |
(undocumented) |
Returns: dict
Keys include:
* "mean": jnp.ndarray (The forecast mean).
* If level is provided, keys for prediction intervals (e.g., "upper_90", "lower_90").
Raises:
* RuntimeError: If fit(y) has not been called.
* ValueError: If h is not positive or exceeds the h the model was trained for.
* ValueError: If level is provided but model.conformal_params is not set.
forecast(self, y, h, X=None, X_future=None, level=None, fitted=False) -> dict
Stateless fit-then-predict on y.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
y |
jnp.ndarray |
- | (undocumented) |
h |
int |
- | (undocumented) |
X |
jnp.ndarray \| None |
None |
(undocumented) |
X_future |
jnp.ndarray \| None |
None |
(undocumented) |
level |
list[int \| float] \| None |
None |
(undocumented) |
fitted |
bool |
False |
(undocumented) |
Returns: dict
Keys include:
* "mean": jnp.ndarray (The forecast mean).
* If level is provided, prediction intervals.
* If fitted=True, "fitted": jnp.ndarray (One-step-ahead fitted values).
Raises:
* NotImplementedError: Exogenous variables (X or X_future) are not supported.
FEDformerForecaster
chronax.models.FEDformerForecaster · inherits FEDformer
Deprecated config-based constructor; prefer :class:FEDformer.
__init__(self, config, *, max_steps=1000, learning_rate=0.0001, batch_size=32, num_lr_decays=3, val_fraction=0.1, val_check_steps=100, early_stop_patience_steps=-1, grad_clip=1.0, weight_decay=0.0, loss='mae', seed=1)
(undocumented)
| Parameter | Type | Default | Description |
|---|---|---|---|
config |
FEDformerConfig |
- | (undocumented) |
max_steps |
int |
1000 |
(undocumented) |
learning_rate |
float |
1e-4 |
(undocumented) |
batch_size |
int |
32 |
(undocumented) |
num_lr_decays |
int |
3 |
(undocumented) |
val_fraction |
float |
0.1 |
(undocumented) |
val_check_steps |
int |
100 |
(undocumented) |
early_stop_patience_steps |
int |
-1 |
(undocumented) |
grad_clip |
float |
1.0 |
(undocumented) |
weight_decay |
float |
0.0 |
(undocumented) |
loss |
Union[str, LossFn, Callable] |
"mae" |
(undocumented) |
seed |
int |
1 |
(undocumented) |