GLU

chronax.tft_layers.GLU

Gated Linear Unit: a * sigmoid(b) where [a, b] = Linear(x) (== torch F.glu).

__init__(self, input_size: int, output_size: int, *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
input_size	int	-	(undocumented)
output_size	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

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

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

GRN

chronax.tft_layers.GRN

Gated Residual Network (NF GRN).

y = MaybeLayerNorm(residual + GLU(lin_i(act(lin_a(a) + lin_c(c))))), where residual = a (or out_proj(a) when output_size is given) and MaybeLayerNorm is identity when output_size == 1. The context c is static (no time axis) and is broadcast across leading non-feature dims.

__init__(self, input_size: int, hidden_size: int, output_size: int | None = None, context_size: int | None = None, dropout: float = 0.0, activation: str = 'ELU', *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
input_size	int	-	(undocumented)
hidden_size	int	-	(undocumented)
output_size	int \| None	None	(undocumented)
context_size	int \| None	None	(undocumented)
dropout	float	0.0	(undocumented)
activation	str	"ELU"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, a: jnp.ndarray, c: jnp.ndarray | None = None, deterministic: bool = True) -> jnp.ndarray

Parameter	Type	Default	Description
a	jnp.ndarray	-	(undocumented)
c	jnp.ndarray \| None	None	(undocumented)
deterministic	bool	True	(undocumented)

VariableSelectionNetwork

chronax.tft_layers.VariableSelectionNetwork

Selects and combines num_inputs embedded variables (NF VSN).

A joint GRN over the flattened variable embeddings (plus optional static context) produces softmax selection weights; each variable is transformed by its own GRN; the outputs are combined by the weights. Works with or without a time axis (static inputs have none).

__init__(self, hidden_size: int, num_inputs: int, dropout: float = 0.0, context_size: int | None = None, activation: str = 'ELU', *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
hidden_size	int	-	(undocumented)
num_inputs	int	-	(undocumented)
dropout	float	0.0	(undocumented)
context_size	int \| None	None	(undocumented)
activation	str	"ELU"	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

__call__(self, x: jnp.ndarray, context: jnp.ndarray | None = None, deterministic: bool = True)

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
context	jnp.ndarray \| None	None	(undocumented)
deterministic	bool	True	(undocumented)
Returns: (out, weights)

ContinuousEmbedding

chronax.tft_layers.ContinuousEmbedding

Per-feature continuous embedding (NF TFTEmbedding, continuous path).

Each scalar feature j maps to x[..., j, None] * vec[j] + bias[j] with learned vec, bias of shape [num_features, hidden]. float32.

__init__(self, num_features: int, hidden_size: int, *, rngs: nnx.Rngs)

Parameter	Type	Default	Description
num_features	int	-	(undocumented)
hidden_size	int	-	(undocumented)
rngs	nnx.Rngs	-	(undocumented)

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

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

InterpretableMultiHeadAttention

chronax.tft_layers.InterpretableMultiHeadAttention

TFT interpretable multi-head attention.

Q and K are multi-head; V is shared across heads (a single value projection), and head outputs are averaged (not concatenated) before the output projection -- this is what makes the attention weights interpretable. Causal: each query attends only to keys at <= its own time index. The mask is computed with jnp.tril (no stored buffer) so the module is vmap/scan-pure.

__init__(self, n_head: int, hidden_size: int, attn_dropout: float = 0.0, *, rngs: nnx.Rngs)

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

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

Parameter	Type	Default	Description
x	jnp.ndarray	-	(undocumented)
deterministic	bool	True	(undocumented)
Returns: (ctx, attn)