make_cell
dilated_rnn_module.make_cell
Build one recurrent cell by NF's cell_type name.
| Parameter | Type | Default | Description |
|---|---|---|---|
| cell_type | str | - | (undocumented) |
| input_size | int | - | (undocumented) |
| hidden_size | int | - | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
Returns: _CellBase
DRNN
dilated_rnn_module.DRNN · inherits nnx.Module
One dilated-RNN group: len(dilations) stacked layers, one per rate.
Layer i consumes layer i-1's output at dilation dilations[i]. NF's
DRNN also returns a list of per-layer tails (inputs[-dilation:]); nothing
downstream reads it, so it is not built here.
__init__(self, *, n_input: int, n_hidden: int, dilations: list[int], cell_type: str = 'GRU', rngs: nnx.Rngs)
| Parameter | Type | Default | Description |
|---|---|---|---|
| n_input | int | - | (undocumented) |
| n_hidden | int | - | (undocumented) |
| dilations | list[int] | - | (undocumented) |
| cell_type | str | 'GRU' | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x: jnp.ndarray) -> jnp.ndarray
x: [B, T, F] (batch-first, as NF's batch_first=True) -> [B, T, H].
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | jnp.ndarray | - | (undocumented) |
Returns: jnp.ndarray
MLPDecoder
dilated_rnn_module.MLPDecoder · inherits nnx.Module
NF common._modules.MLP with ReLU and dropout=0.0.
num_layers=1 is a bare Linear(in -> out) with no activation;
num_layers=n >= 2 is Linear(in -> hidden), then n - 2 hidden
Linear(hidden -> hidden) blocks, then Linear(hidden -> out), with ReLU
after every layer except the last.
__init__(self, *, in_features: int, hidden_size: int, out_features: int, num_layers: int, rngs: nnx.Rngs)
| Parameter | Type | Default | Description |
|---|---|---|---|
| in_features | int | - | (undocumented) |
| hidden_size | int | - | (undocumented) |
| out_features | int | - | (undocumented) |
| num_layers | int | - | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x: jnp.ndarray) -> jnp.ndarray
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | jnp.ndarray | - | (undocumented) |
Returns: jnp.ndarray
DilatedRNNNet
dilated_rnn_module.DilatedRNNNet · inherits nnx.Module
Full DilatedRNN backbone: DRNN stack -> context adapter -> MLP decoder.
I/O mirrors the other Chronax neural nets: __call__(x: [B, L, 1]) -> [B, h, 1].
Written N-generically in the feature dimension (in_features), so a future
exogenous path can widen the encoder input without touching the rest.
__init__(self, *, h: int, input_size: int, in_features: int = 1, cell_type: str = 'LSTM', dilations: list[list[int]] \| None = None, encoder_hidden_size: int = 128, decoder_hidden_size: int = 128, decoder_layers: int = 2, rngs: nnx.Rngs)
| Parameter | Type | Default | Description |
|---|---|---|---|
| h | int | - | (undocumented) |
| input_size | int | - | (undocumented) |
| in_features | int | 1 | (undocumented) |
| cell_type | str | 'LSTM' | (undocumented) |
| dilations | list[list[int]] | None | None | (undocumented) |
| encoder_hidden_size | int | 128 | (undocumented) |
| decoder_hidden_size | int | 128 | (undocumented) |
| decoder_layers | int | 2 | (undocumented) |
| rngs | nnx.Rngs | - | (undocumented) |
__call__(self, x: jnp.ndarray, deterministic: bool = True) -> jnp.ndarray
x: [B, L, in_features] -> [B, h, 1].
deterministic is accepted for interface symmetry with the other neural
ports; DilatedRNN has no stochastic layers (NF builds its decoder with
dropout=0.0 and its cells never dropout at one layer per RNN), so the
flag changes nothing.
Parameters:
| Parameter | Type | Default | Description |
|---|---|---|---|
| x | jnp.ndarray | - | (undocumented) |
| deterministic | bool | True | (undocumented) |
Returns: jnp.ndarray