TokenEmbedding
chronax.models.vanillatransformer_module.TokenEmbedding · inherits nnx.Module
Circular k=3 Conv mapping the univariate channel to hidden_size.
NF TokenEmbedding is Conv1d(c_in=1, hidden, kernel_size=3, padding=1, padding_mode='circular', bias=False) over [B, C, L]. flax nnx.Conv is channels-last, so the input [B, L, 1] is consumed directly with padding='CIRCULAR'; no permute needed.
__init__(self, *, hidden_size: int, rngs: nnx.Rngs)
Initializes the TokenEmbedding module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| hidden_size | int | - | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x: jnp.ndarray) -> jnp.ndarray
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | jnp.ndarray | - | (undocumented) |
DataEmbedding
chronax.models.vanillatransformer_module.DataEmbedding · inherits nnx.Module
Token conv + fixed sinusoidal positional embedding + dropout.
Mirrors NF DataEmbedding for the univariate, no-exogenous case (temporal_embedding omitted; pos_embedding=True).
__init__(self, *, hidden_size: int, dropout: float, rngs: nnx.Rngs)
Initializes the DataEmbedding module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| hidden_size | int | - | (undocumented) |
| dropout | float | - | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray
x: [B, T, 1] -> [B, T, hidden].
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | jnp.ndarray | - | (undocumented) |
| deterministic | bool | - | (undocumented) |
AttentionLayer
chronax.models.vanillatransformer_module.AttentionLayer · inherits nnx.Module
Multi-head full (softmax dot-product) attention, self- or cross-.
Folds NF AttentionLayer + FullAttention (output_attention=False, no mask) into one module. q is projected from q_in; k/v from kv_in (equal for self-attention). Scale d_k**-0.5; dropout on the attention weights; torch nn.Linear init.
__init__(self, *, hidden_size: int, n_heads: int, attn_dropout: float, rngs: nnx.Rngs)
Initializes the AttentionLayer module.
| 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)
q_in: [B, Lq, hidden], kv_in: [B, Lk, hidden] -> [B, Lq, hidden].
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| q_in | - | - | (undocumented) |
| kv_in | - | - | (undocumented) |
| deterministic | bool | - | (undocumented) |
TransEncoderLayer
chronax.models.vanillatransformer_module.TransEncoderLayer · inherits nnx.Module
One encoder layer (post-norm), matching NF TransEncoderLayer.
new_x = attn(x, x); x = x + drop(new_x); y = x = norm1(x); y = drop(act(conv1(y))); y = drop(conv2(y)); return norm2(x + y). The two Conv1d(kernel_size=1) are pointwise linears with bias.
__init__(self, *, hidden_size, n_heads, conv_hidden_size, dropout, activation="gelu", rngs: nnx.Rngs)
Initializes the TransEncoderLayer module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| hidden_size | - | - | (undocumented) |
| n_heads | - | - | (undocumented) |
| conv_hidden_size | - | - | (undocumented) |
| dropout | - | - | (undocumented) |
| activation | str | "gelu" | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x, deterministic: bool)
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | - | - | (undocumented) |
| deterministic | bool | - | (undocumented) |
TransEncoder
chronax.models.vanillatransformer_module.TransEncoder · inherits nnx.Module
Stack of encoder_layers encoder layers + final LayerNorm.
__init__(self, *, encoder_layers, hidden_size, n_heads, conv_hidden_size, dropout, activation="gelu", rngs: nnx.Rngs)
Initializes the TransEncoder module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| encoder_layers | - | - | (undocumented) |
| hidden_size | - | - | (undocumented) |
| n_heads | - | - | (undocumented) |
| conv_hidden_size | - | - | (undocumented) |
| dropout | - | - | (undocumented) |
| activation | str | "gelu" | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x, deterministic: bool)
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | - | - | (undocumented) |
| deterministic | bool | - | (undocumented) |
TransDecoderLayer
chronax.models.vanillatransformer_module.TransDecoderLayer · inherits nnx.Module
One decoder layer, matching NF TransDecoderLayer (no causal mask).
x = x + drop(self_attn(x, x)); x = norm1(x); x = x + drop(cross_attn(x, cross)); y = x = norm2(x); y = drop(act(conv1(y))); y = drop(conv2(y)); return norm3(x + y).
__init__(self, *, hidden_size, n_heads, conv_hidden_size, dropout, activation="gelu", rngs: nnx.Rngs)
Initializes the TransDecoderLayer module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| hidden_size | - | - | (undocumented) |
| n_heads | - | - | (undocumented) |
| conv_hidden_size | - | - | (undocumented) |
| dropout | - | - | (undocumented) |
| activation | str | "gelu" | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x, cross, deterministic: bool)
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | - | - | (undocumented) |
| cross | - | - | (undocumented) |
| deterministic | bool | - | (undocumented) |
TransDecoder
chronax.models.vanillatransformer_module.TransDecoder · inherits nnx.Module
Stack of decoder_layers decoder layers + LayerNorm + projection.
__init__(self, *, decoder_layers, hidden_size, n_heads, conv_hidden_size, dropout, activation="gelu", c_out=1, rngs: nnx.Rngs)
Initializes the TransDecoder module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| decoder_layers | - | - | (undocumented) |
| hidden_size | - | - | (undocumented) |
| n_heads | - | - | (undocumented) |
| conv_hidden_size | - | - | (undocumented) |
| dropout | - | - | (undocumented) |
| activation | str | "gelu" | (undocumented) |
| c_out | int | 1 | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x, cross, deterministic: bool)
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | - | - | (undocumented) |
| cross | - | - | (undocumented) |
| deterministic | bool | - | (undocumented) |
VanillaTransformerNet
chronax.models.vanillatransformer_module.VanillaTransformerNet · inherits nnx.Module
Full encoder-decoder backbone: [B, input_size, 1] -> [B, h, 1].
Identity scaling (NF default): operates in raw scale. During training pass deterministic=False; nnx.Rngs supplies the dropout stream.
__init__(self, *, h, input_size, hidden_size, n_heads, conv_hidden_size, encoder_layers, decoder_layers, dropout, activation="gelu", decoder_input_size_multiplier=0.5, rngs: nnx.Rngs)
Initializes the VanillaTransformerNet module.
| Parameter | Type | Default | Description |
|---|---|---|---|
| h | - | - | (undocumented) |
| input_size | - | - | (undocumented) |
| hidden_size | - | - | (undocumented) |
| n_heads | - | - | (undocumented) |
| conv_hidden_size | - | - | (undocumented) |
| encoder_layers | - | - | (undocumented) |
| decoder_layers | - | - | (undocumented) |
| dropout | - | - | (undocumented) |
| activation | str | "gelu" | (undocumented) |
| decoder_input_size_multiplier | float | 0.5 | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray
(undocumented)
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | jnp.ndarray | - | (undocumented) |
| deterministic | bool | - | (undocumented) |