StaticCovariateEncoder

tft_module.StaticCovariateEncoder

Encodes static covariates into the four TFT context vectors. A VSN over the static variables yields a context, which separate GRNs map to cs (variable-selection context), ce (enrichment context), and the per-layer LSTM initial hidden/cell states ch/cc (cc == ch for GRU).

__init__(self, hidden_size, num_static, dropout, activation, rnn_type="lstm", n_rnn_layers=1, *, rngs: nnx.Rngs)

(No prose summary)

Parameter	Type	Default	Description
hidden_size	-	-	(undocumented)
num_static	-	-	(undocumented)
dropout	-	-	(undocumented)
activation	-	-	(undocumented)
rnn_type	str	"lstm"	(undocumented)
n_rnn_layers	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, s: jnp.ndarray, deterministic: bool = True)

(No prose summary)

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

Returns: tuple (The four context vectors and VSN weights: cs, ce, ch, cc, weights).

TemporalCovariateEncoder

tft_module.TemporalCovariateEncoder

Locality-enhancement seq2seq: VSN-gated history/future + LSTM/GRU encoder-decoder. History and future variables are variable-selected (with static context cs), the encoder recurrence runs over history from the static initial carry (cc, ch), the decoder recurrence runs over the future continuing from the encoder's final carry, and the concatenated outputs are GLU-gated with a residual + LayerNorm (NF input_gate).

__init__(self, hidden_size, num_hist_vars, num_futr_vars, dropout, activation, rnn_type="lstm", n_rnn_layers=1, *, rngs: nnx.Rngs)

(No prose summary)

Parameter	Type	Default	Description
hidden_size	-	-	(undocumented)
num_hist_vars	-	-	(undocumented)
num_futr_vars	-	-	(undocumented)
dropout	-	-	(undocumented)
activation	-	-	(undocumented)
rnn_type	str	"lstm"	(undocumented)
n_rnn_layers	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, hist, futr, cs, ch, cc, deterministic=True)

(No prose summary)

Parameter	Type	Default	Description
hist	-	-	(undocumented)
futr	-	-	(undocumented)
cs	-	-	(undocumented)
ch	-	-	(undocumented)
cc	-	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: tuple (The combined temporal features and VSN weights: temporal, hist_w, futr_w).

TemporalFusionDecoder

tft_module.TemporalFusionDecoder

Static enrichment -> masked interpretable attention -> position-wise GRN. Each block is gated (GLU) with a residual + LayerNorm. After attention the encoder steps are discarded, keeping only the h horizon steps (NF slices [input_size:]).

__init__(self, n_head, hidden_size, dropout, attn_dropout, activation, *, rngs: nnx.Rngs)

(No prose summary)

Parameter	Type	Default	Description
n_head	-	-	(undocumented)
hidden_size	-	-	(undocumented)
dropout	-	-	(undocumented)
attn_dropout	-	-	(undocumented)
activation	-	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, temporal, ce, input_size, deterministic=True)

(No prose summary)

Parameter	Type	Default	Description
temporal	-	-	(undocumented)
ce	-	-	(undocumented)
input_size	-	-	(undocumented)
deterministic	bool	True	(undocumented)

Returns: tuple (The decoded features and attention weights: x, attn_w).

TFTNet

tft_module.TFTNet

Full TFT backbone: embed -> static contexts -> seq2seq -> fusion decoder -> head. I/O mirrors NF TFT.forward but with explicit kwargs instead of a windows_batch dict. __call__ returns [B, h, outputsize_multiplier].

__init__(self, *, h, input_size, hidden_size=128, n_head=4, attn_dropout=0.0, dropout=0.1, grn_activation="ELU", rnn_type="lstm", n_rnn_layers=1, stat_exog_size=0, hist_exog_size=0, futr_exog_size=0, tgt_size=1, outputsize_multiplier=1, rngs: nnx.Rngs)

(No prose summary)

Parameter	Type	Default	Description
h	-	-	(undocumented)
input_size	-	-	(undocumented)
hidden_size	int	128	(undocumented)
n_head	int	4	(undocumented)
attn_dropout	float	0.0	(undocumented)
dropout	float	0.1	(undocumented)
grn_activation	str	"ELU"	(undocumented)
rnn_type	str	"lstm"	(undocumented)
n_rnn_layers	int	1	(undocumented)
stat_exog_size	int	0	(undocumented)
hist_exog_size	int	0	(undocumented)
futr_exog_size	int	0	(undocumented)
tgt_size	int	1	(undocumented)
outputsize_multiplier	int	1	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, insample_y, hist_exog=None, futr_exog=None, stat_exog=None, deterministic=True)

(No prose summary)

Parameter	Type	Default	Description
insample_y	-	-	(undocumented)
hist_exog	-	None	(undocumented)
futr_exog	-	None	(undocumented)
stat_exog	-	None	(undocumented)
deterministic	bool	True	(undocumented)

Returns: jnp.ndarray (Shape [B, h, outputsize_multiplier]).