Recurrent Layers

Recurrent cells and layers.

class taktiny.nn.RNN(input_size, hidden_size, num_layers=1, *, nonlinearity='tanh', bias=True, batch_first=False, dropout=0.0, bidirectional=False, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, recurrent_initializer=None, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None, unroll=1)[source]

Bases: RecurrentBase

Apply a stacked Elman RNN to a complete sequence.

At each time step and layer, the recurrence is

\[h_t = \phi\left( x_t W_{ih} + b_{ih} + h_{t-1} W_{hh} + b_{hh} \right).\]

Inputs are [sequence, input] when unbatched. Batched inputs are [sequence, batch, input] unless batch_first=True. The returned hidden state is ordered by layer, then direction. Bidirectional outputs concatenate the forward and reverse hidden states on the feature axis.

Parameters:

  • input_size (int) – Number of features in each input time step.

  • hidden_size (int) – Number of features in each directional hidden state.

  • num_layers (int) – Number of stacked recurrent layers. Defaults to 1.

  • nonlinearity (Union[Literal['tanh', 'relu'], Callable[[Array], Array]]) – 'tanh', 'relu', or a callable activation. Defaults to 'tanh'.

  • bias (bool) – Whether both cell projections include biases. Defaults to True.

  • batch_first (bool) – Whether batched inputs use batch-major layout. Defaults to False.

  • dropout (float) – Dropout probability between layers. Dropout is active only during training and is not applied after the last layer. Defaults to 0.0.

  • bidirectional (bool) – Whether each layer scans in both directions. Defaults to False.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to parameter initializers.

  • rngs (Rngs) – Random number generator used to initialize all cells.

  • kernel_initializer (Initializer) – Input-kernel initializer. Defaults to LeCun uniform initialization.

  • recurrent_initializer (Initializer | None) – Recurrent-kernel initializer. Defaults to kernel_initializer.

  • bias_initializer (Initializer) – Bias initializer. Defaults to zeros.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for projection kernels.

  • dot_general (DotGeneral | None) – Optional dot_general implementation used by cells.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for the input and hidden axes.

  • partition_spec (P | None) – Optional partition specifications for the input and hidden axes.

  • kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to kernel parameters.

  • bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to bias parameters.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to cell projections.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to cell projections.

  • unroll (int | bool) – Loop-unrolling option passed to jax.lax.scan(). Defaults to 1.

References

Examples

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> rnn = nn.RNN(3, 4, rngs=nn.Rngs(0))
>>> output, hidden = rnn(jnp.ones((5, 3)))
>>> (output.shape, hidden.shape)
((5, 4), (1, 4))
class taktiny.nn.RNNCell(input_size, hidden_size, *, nonlinearity='tanh', bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, recurrent_initializer=None, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]

Bases: Module

Apply one Elman recurrent-network step.

\[h_t = \phi\left( x_t W_{ih} + b_{ih} + h_{t-1} W_{hh} + b_{hh} \right)\]

Here, \(\phi\) is tanh, relu, or a custom activation. The output is \(h_t\). State is stored as the one-element tuple (hidden,) so the cell follows the (state, input) convention used by jax.lax.scan(). Both bias terms are omitted when bias=False.

axis_names and partition_spec describe the semantic input and hidden dimensions. The recurrent kernel leaves its contracting dimension unassigned, preventing the same logical mesh axis from appearing twice.

Parameters:

  • input_size (int) – Number of features in \(x_t\).

  • hidden_size (int) – Number of features in \(h_t\).

  • nonlinearity (Union[Literal['tanh', 'relu'], Callable[[Array], Array]]) – Hidden-state activation: 'tanh', 'relu', or a callable. Defaults to 'tanh'.

  • bias (bool) – Whether both projections include a learnable bias. Defaults to True.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to parameter initializers. Defaults to each initializer’s default data type.

  • rngs (Rngs) – Random number generator used to initialize parameters.

  • kernel_initializer (Initializer) – Input-kernel initializer. Defaults to LeCun uniform initialization.

  • recurrent_initializer (Initializer | None) – Recurrent-kernel initializer. Defaults to kernel_initializer.

  • bias_initializer (Initializer) – Bias initializer. Defaults to zeros.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for both kernels.

  • dot_general (DotGeneral | None) – Optional dot_general implementation used by the projections.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for the input and hidden axes.

  • partition_spec (P | None) – Optional partition specifications for the input and hidden axes.

  • kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to both kernel parameters.

  • bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to both bias parameters.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to the projections.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to the projections.

Examples

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> cell = nn.RNNCell(3, 4, rngs=nn.Rngs(0))
>>> state = cell.initial_state((2,))
>>> next_state, output = cell(state, jnp.ones((2, 3)))
>>> (next_state[0].shape, output.shape)
((2, 4), (2, 4))
class taktiny.nn.LSTM(input_size, hidden_size, num_layers=1, *, bias=True, batch_first=False, dropout=0.0, bidirectional=False, proj_size=0, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, recurrent_initializer=None, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None, unroll=1)[source]

Bases: RecurrentBase

Apply a stacked long short-term memory network to a sequence.

Each cell computes input, forget, candidate, and output gates from the current input and previous hidden state, then updates

\[\begin{split}\begin{aligned} c_t &= f_t \odot c_{t-1} + i_t \odot g_t \\ \widetilde{h}_t &= o_t \odot \tanh(c_t) \end{aligned}\end{split}\]

Without projection, \(h_t = \widetilde{h}_t\). With proj_size > 0, \(h_t = \widetilde{h}_t W_{hr}\).

The returned state is (hidden, cell). With proj_size > 0, hidden outputs have width proj_size while cell states retain hidden_size. Bidirectional outputs concatenate both directions on the feature axis.

Parameters:

  • input_size (int) – Number of features in each input time step.

  • hidden_size (int) – Number of features in each cell state.

  • num_layers (int) – Number of stacked recurrent layers. Defaults to 1.

  • bias (bool) – Whether the input and recurrent gate projections include biases. Defaults to True.

  • batch_first (bool) – Whether batched inputs use batch-major layout. Defaults to False.

  • dropout (float) – Dropout probability between layers. Dropout is active only during training and is not applied after the last layer. Defaults to 0.0.

  • bidirectional (bool) – Whether each layer scans in both directions. Defaults to False.

  • proj_size (int) – Per-direction hidden output width. 0 disables the projection; otherwise it must be smaller than hidden_size.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to parameter initializers.

  • rngs (Rngs) – Random number generator used to initialize all cells.

  • kernel_initializer (Initializer) – Input-kernel initializer. Defaults to LeCun uniform initialization.

  • recurrent_initializer (Initializer | None) – Recurrent and output-projection initializer. Defaults to kernel_initializer.

  • bias_initializer (Initializer) – Gate-bias initializer. Defaults to zeros.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for projection kernels.

  • dot_general (DotGeneral | None) – Optional dot_general implementation used by cells.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for the input and hidden axes.

  • partition_spec (P | None) – Optional partition specifications for the input and hidden axes.

  • kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to kernel parameters.

  • bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to gate-bias parameters.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to cell projections.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to cell projections.

  • unroll (int | bool) – Loop-unrolling option passed to jax.lax.scan(). Defaults to 1.

References

Examples

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> lstm = nn.LSTM(3, 5, proj_size=2, rngs=nn.Rngs(0))
>>> output, (hidden, cell) = lstm(jnp.ones((4, 3)))
>>> (output.shape, hidden.shape, cell.shape)
((4, 2), (1, 2), (1, 5))
class taktiny.nn.LSTMCell(input_size, hidden_size, *, proj_size=0, bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, recurrent_initializer=None, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]

Bases: Module

Apply one long short-term memory step.

\[\begin{split}\begin{aligned} i_t &= \sigma(x_t W_{ii} + b_{ii} + h_{t-1} W_{hi} + b_{hi}) \\ f_t &= \sigma(x_t W_{if} + b_{if} + h_{t-1} W_{hf} + b_{hf}) \\ g_t &= \tanh(x_t W_{ig} + b_{ig} + h_{t-1} W_{hg} + b_{hg}) \\ o_t &= \sigma(x_t W_{io} + b_{io} + h_{t-1} W_{ho} + b_{ho}) \\ c_t &= f_t \odot c_{t-1} + i_t \odot g_t \\ \widetilde{h}_t &= o_t \odot \tanh(c_t) \end{aligned}\end{split}\]

Without projection, \(h_t = \widetilde{h}_t\). When proj_size is nonzero, the exposed hidden state is instead

\[h_t = \widetilde{h}_t W_{hr}.\]

State is (hidden, cell). A projected hidden state has width proj_size, while the cell state always retains hidden_size. Both bias terms in each gate are omitted when bias=False.

Parameters:

  • input_size (int) – Number of features in \(x_t\).

  • hidden_size (int) – Number of features in the cell state \(c_t\).

  • proj_size (int) – Width of the exposed hidden state. 0 disables the projection. A nonzero value must be smaller than hidden_size.

  • bias (bool) – Whether the input and recurrent gate projections include learnable biases. Defaults to True.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to parameter initializers. Defaults to each initializer’s default data type.

  • rngs (Rngs) – Random number generator used to initialize parameters.

  • kernel_initializer (Initializer) – Input-kernel initializer. Defaults to LeCun uniform initialization.

  • recurrent_initializer (Initializer | None) – Recurrent and output-projection initializer. Defaults to kernel_initializer.

  • bias_initializer (Initializer) – Gate-bias initializer. Defaults to zeros.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for all kernels.

  • dot_general (DotGeneral | None) – Optional dot_general implementation used by the projections.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for the input and hidden axes.

  • partition_spec (P | None) – Optional partition specifications for the input and hidden axes.

  • kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to kernel parameters.

  • bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to gate-bias parameters.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to the projections.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to the projections.

Examples

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> cell = nn.LSTMCell(3, 5, proj_size=2, rngs=nn.Rngs(0))
>>> state = cell.initial_state((4,))
>>> next_state, output = cell(state, jnp.ones((4, 3)))
>>> (output.shape, next_state[1].shape)
((4, 2), (4, 5))
class taktiny.nn.GRU(input_size, hidden_size, num_layers=1, *, bias=True, batch_first=False, dropout=0.0, bidirectional=False, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, recurrent_initializer=None, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None, unroll=1)[source]

Bases: RecurrentBase

Apply a stacked gated recurrent unit network to a sequence.

Each cell combines a reset gate \(r_t\), update gate \(z_t\), and candidate activation \(n_t\):

\[\begin{split}\begin{aligned} r_t &= \sigma(x_t W_{ir} + b_{ir} + h_{t-1} W_{hr} + b_{hr}) \\ z_t &= \sigma(x_t W_{iz} + b_{iz} + h_{t-1} W_{hz} + b_{hz}) \\ n_t &= \tanh\left(x_t W_{in} + b_{in} + r_t \odot (h_{t-1} W_{hn} + b_{hn})\right) \\ h_t &= (1-z_t) \odot n_t + z_t \odot h_{t-1} \end{aligned}\end{split}\]

This implementation uses the reset-after GRU form. Bidirectional outputs concatenate the forward and reverse hidden states on the feature axis.

Parameters:

  • input_size (int) – Number of features in each input time step.

  • hidden_size (int) – Number of features in each directional hidden state.

  • num_layers (int) – Number of stacked recurrent layers. Defaults to 1.

  • bias (bool) – Whether the input and recurrent gate projections include biases. Defaults to True.

  • batch_first (bool) – Whether batched inputs use batch-major layout. Defaults to False.

  • dropout (float) – Dropout probability between layers. Dropout is active only during training and is not applied after the last layer. Defaults to 0.0.

  • bidirectional (bool) – Whether each layer scans in both directions. Defaults to False.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to parameter initializers.

  • rngs (Rngs) – Random number generator used to initialize all cells.

  • kernel_initializer (Initializer) – Input-kernel initializer. Defaults to LeCun uniform initialization.

  • recurrent_initializer (Initializer | None) – Recurrent-kernel initializer. Defaults to kernel_initializer.

  • bias_initializer (Initializer) – Gate-bias initializer. Defaults to zeros.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for projection kernels.

  • dot_general (DotGeneral | None) – Optional dot_general implementation used by cells.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for the input and hidden axes.

  • partition_spec (P | None) – Optional partition specifications for the input and hidden axes.

  • kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to kernel parameters.

  • bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to gate-bias parameters.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to cell projections.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to cell projections.

  • unroll (int | bool) – Loop-unrolling option passed to jax.lax.scan(). Defaults to 1.

Examples

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> gru = nn.GRU(3, 4, rngs=nn.Rngs(0))
>>> output, hidden = gru(jnp.ones((5, 3)))
>>> (output.shape, hidden.shape)
((5, 4), (1, 4))
class taktiny.nn.GRUCell(input_size, hidden_size, *, bias=True, dtype=None, rngs, kernel_initializer=<function variance_scaling.<locals>.init>, recurrent_initializer=None, bias_initializer=<function zeros>, quant=None, dot_general=None, axis_names=None, partition_spec=None, kernel_metadata=None, bias_metadata=None, precision=None, preferred_element_type=None)[source]

Bases: Module

Apply one gated recurrent unit step.

\[\begin{split}\begin{aligned} r_t &= \sigma(x_t W_{ir} + b_{ir} + h_{t-1} W_{hr} + b_{hr}) \\ z_t &= \sigma(x_t W_{iz} + b_{iz} + h_{t-1} W_{hz} + b_{hz}) \\ n_t &= \tanh\left(x_t W_{in} + b_{in} + r_t \odot (h_{t-1} W_{hn} + b_{hn})\right) \\ h_t &= (1-z_t) \odot n_t + z_t \odot h_{t-1} \end{aligned}\end{split}\]

Here, \(r_t\), \(z_t\), and \(n_t\) are the reset, update, and candidate activations. This implementation uses the reset-after form: the reset gate is applied after the candidate’s recurrent projection. State is stored as (hidden,), and the output is the new hidden state.

Parameters:

  • input_size (int) – Number of features in \(x_t\).

  • hidden_size (int) – Number of features in \(h_t\).

  • bias (bool) – Whether the input and recurrent gate projections include learnable biases. Defaults to True.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Data type passed to parameter initializers. Defaults to each initializer’s default data type.

  • rngs (Rngs) – Random number generator used to initialize parameters.

  • kernel_initializer (Initializer) – Input-kernel initializer. Defaults to LeCun uniform initialization.

  • recurrent_initializer (Initializer | None) – Recurrent-kernel initializer. Defaults to kernel_initializer.

  • bias_initializer (Initializer) – Gate-bias initializer. Defaults to zeros.

  • quant (str | QuantizationRule | PtqProvider | Sequence[QuantizationRule] | None) – Optional Qwix quantization configuration for both kernels.

  • dot_general (DotGeneral | None) – Optional dot_general implementation used by the projections.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for the input and hidden axes.

  • partition_spec (P | None) – Optional partition specifications for the input and hidden axes.

  • kernel_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to both kernel parameters.

  • bias_metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata attached to both bias parameters.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision], DotAlgorithm, DotAlgorithmPreset]) – Dot-product precision forwarded to the projections.

  • preferred_element_type (Union[str, type[Any], dtype, SupportsDType, None]) – Preferred result and accumulation data type forwarded to the projections.

Examples

>>> import jax.numpy as jnp
>>> from taktiny import nn
>>> cell = nn.GRUCell(3, 4, rngs=nn.Rngs(0))
>>> state = cell.initial_state((2,))
>>> next_state, output = cell(state, jnp.ones((2, 3)))
>>> (next_state[0].shape, output.shape)
((2, 4), (2, 4))
class taktiny.nn.RecurrentBase(input_size, hidden_size, num_layers, *, bias, batch_first, dropout, bidirectional, dtype, axis_names, unroll, output_size, state_sizes, cell_factory)[source]

Bases: Module

Provide stacked sequence traversal for recurrent network subclasses.

Each layer is evaluated with jax.lax.scan(). A bidirectional layer scans independently in both directions and concatenates their outputs on the feature axis. Dropout is applied between stacked layers only, never after the final layer.

This is the shared implementation behind RNN, LSTM, and GRU; users normally instantiate one of those subclasses.

Parameters:

  • input_size (int) – Number of features in each input time step.

  • hidden_size (int) – Internal hidden width reported by the subclass.

  • num_layers (int) – Number of stacked recurrent layers.

  • bias (bool) – Whether recurrent cells use biases.

  • batch_first (bool) – Whether batched inputs use (batch, sequence, feature) instead of (sequence, batch, feature).

  • dropout (float) – Dropout probability between recurrent layers.

  • bidirectional (bool) – Whether every layer scans in both directions.

  • dtype (Union[str, type[Any], dtype, SupportsDType] | None) – Parameter data type requested by the subclass.

  • axis_names (tuple[str | None, ...] | None) – Optional logical names for input and hidden axes.

  • unroll (int | bool) – Loop-unrolling option forwarded to jax.lax.scan().

  • output_size (int) – Per-direction output width of one cell.

  • state_sizes (tuple[int, ...]) – Trailing width of each state component.

  • cell_factory (Callable[[int, str | None], Module]) – Factory receiving a layer input width and its logical axis name, and returning a recurrent cell.

Variables:

  • layers – Recurrent layers in stack order.

  • num_directions – 2 for a bidirectional network, otherwise 1.

  • output_size – Per-direction output width.

  • state_sizes – Trailing width of each recurrent state component.

class taktiny.nn.RecurrentLayer(forward_cell, reverse_cell=None)[source]

Bases: Module

Hold the forward and optional reverse cells for one recurrent layer.

Parameters:

  • forward_cell (Module) – Cell applied in sequence order.

  • reverse_cell (Module | None) – Optional cell applied in reverse sequence order. If present, the layer is bidirectional.

Variables:

  • forward_cell – The forward recurrent cell.

  • reverse_cell – The reverse recurrent cell, or None.