From 18d48773690c33b307fa48d9ea0f50c309269c33 Mon Sep 17 00:00:00 2001 From: MartinuzziFrancesco Date: Wed, 24 Jun 2026 14:56:30 +0200 Subject: [PATCH 1/2] feat: add taugrucell --- docs/api/cells.rst | 1 + docs/api/layers.rst | 1 + docs/generated/torchrecurrent.tauGRU.rst | 8 + docs/generated/torchrecurrent.tauGRUCell.rst | 8 + docs/models.rst | 3 + tests/test_cells.py | 31 ++ tests/test_layers.py | 3 + torchrecurrent/__init__.py | 4 + torchrecurrent/cells/__init__.py | 3 + torchrecurrent/cells/taugru_cell.py | 314 +++++++++++++++++++ 10 files changed, 376 insertions(+) create mode 100644 docs/generated/torchrecurrent.tauGRU.rst create mode 100644 docs/generated/torchrecurrent.tauGRUCell.rst create mode 100644 torchrecurrent/cells/taugru_cell.py diff --git a/docs/api/cells.rst b/docs/api/cells.rst index 7a8d4bd..6402afd 100755 --- a/docs/api/cells.rst +++ b/docs/api/cells.rst @@ -38,6 +38,7 @@ This page documents all custom recurrent cells provided in the `torchrecurrent.c torchrecurrent.SGUCell torchrecurrent.SGRNCell torchrecurrent.STARCell + torchrecurrent.tauGRUCell torchrecurrent.UGRNNCell torchrecurrent.UnICORNNCell torchrecurrent.WMCLSTMCell diff --git a/docs/api/layers.rst b/docs/api/layers.rst index 024f70c..2e6487c 100644 --- a/docs/api/layers.rst +++ b/docs/api/layers.rst @@ -38,6 +38,7 @@ This page documents all custom recurrent layers provided in the `torchrecurrent` torchrecurrent.SGU torchrecurrent.SGRN torchrecurrent.STAR + torchrecurrent.tauGRU torchrecurrent.UGRNN torchrecurrent.UnICORNN torchrecurrent.WMCLSTM diff --git a/docs/generated/torchrecurrent.tauGRU.rst b/docs/generated/torchrecurrent.tauGRU.rst new file mode 100644 index 0000000..3b2a1bf --- /dev/null +++ b/docs/generated/torchrecurrent.tauGRU.rst @@ -0,0 +1,8 @@ +torchrecurrent.tauGRU +===================== + +.. currentmodule:: torchrecurrent + +.. autoclass:: tauGRU + :members: + :show-inheritance: diff --git a/docs/generated/torchrecurrent.tauGRUCell.rst b/docs/generated/torchrecurrent.tauGRUCell.rst new file mode 100644 index 0000000..f2a7b50 --- /dev/null +++ b/docs/generated/torchrecurrent.tauGRUCell.rst @@ -0,0 +1,8 @@ +torchrecurrent.tauGRUCell +========================= + +.. currentmodule:: torchrecurrent + +.. autoclass:: tauGRUCell + :members: + :show-inheritance: diff --git a/docs/models.rst b/docs/models.rst index 5f4114e..0a3daaa 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -110,6 +110,9 @@ references and official implementations where available. * - :doc:`STAR ` - `TPAMI 2022 `__ - `0zgur0/STAckable-Recurrent-network `__ + * - :doc:`tauGRU ` + - `AISTATS 2025 `__ + - – * - :doc:`UGRNN ` - `ICLR 2017 `__ - – diff --git a/tests/test_cells.py b/tests/test_cells.py index b2fff8a..17a0987 100755 --- a/tests/test_cells.py +++ b/tests/test_cells.py @@ -38,6 +38,7 @@ SGUCell, SGRNCell, STARCell, + tauGRUCell, UGRNNCell, UnICORNNCell, WMCLSTMCell, @@ -74,6 +75,7 @@ (SGUCell, 3, 5, False), (SGRNCell, 3, 5, False), (STARCell, 3, 5, False), + (tauGRUCell, 3, 5, False), (UGRNNCell, 3, 5, False), (UnICORNNCell, 3, 5, True), (WMCLSTMCell, 3, 5, True), @@ -132,6 +134,35 @@ def test_reslstm_cell_parameter_shapes(): assert cell.weight_ph.shape == (27,) +def test_taugru_cell_parameter_shapes(): + cell = tauGRUCell(4, 9) + + assert cell.weight_ih.shape == (36, 4) + assert cell.weight_hh.shape == (36, 9) + assert cell.bias_ih.shape == (36,) + assert cell.bias_hh.shape == (36,) + + +def test_taugru_cell_uses_delayed_state(): + cell = tauGRUCell(1, 1) + with torch.no_grad(): + cell.weight_ih.zero_() + cell.weight_hh.zero_() + cell.bias_ih.zero_() + cell.bias_hh.zero_() + cell.weight_hh[1, 0] = 1.0 + cell.bias_hh[2] = 20.0 + cell.bias_hh[3] = 20.0 + + x = torch.zeros(1, 1) + h = torch.zeros(1, 1) + delayed = torch.ones(1, 1) + + out = cell(x, h, delayed) + + assert torch.allclose(out, torch.tanh(delayed), atol=1e-4) + + @pytest.mark.parametrize("Cell, in_size, hid_size, _", CELL_CASES) def test_cell_gradients(Cell, in_size, hid_size, _): """A quick smoke test: outputs should be differentiable wrt parameters.""" diff --git a/tests/test_layers.py b/tests/test_layers.py index c742998..bee0445 100755 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -30,6 +30,7 @@ SGU, SGRN, STAR, + tauGRU, UGRNN, UnICORNN, WMCLSTM, @@ -63,6 +64,7 @@ SGU, SGRN, STAR, + tauGRU, UGRNN, UnICORNN, WMCLSTM, @@ -97,6 +99,7 @@ (SGU, False), (SGRN, False), (STAR, False), + (tauGRU, False), (UGRNN, False), (UnICORNN, True), (WMCLSTM, True), diff --git a/torchrecurrent/__init__.py b/torchrecurrent/__init__.py index 4dec564..eecd4e5 100755 --- a/torchrecurrent/__init__.py +++ b/torchrecurrent/__init__.py @@ -33,6 +33,7 @@ SGUCell, SGRNCell, STARCell, + tauGRUCell, UGRNNCell, UnICORNNCell, WMCLSTMCell, @@ -71,6 +72,7 @@ SGU, SGRN, STAR, + tauGRU, UGRNN, UnICORNN, WMCLSTM, @@ -139,6 +141,8 @@ "SGRNCell", "STAR", "STARCell", + "tauGRU", + "tauGRUCell", "UGRNN", "UGRNNCell", "UnICORNN", diff --git a/torchrecurrent/cells/__init__.py b/torchrecurrent/cells/__init__.py index 53bc57a..7d0aa18 100755 --- a/torchrecurrent/cells/__init__.py +++ b/torchrecurrent/cells/__init__.py @@ -29,6 +29,7 @@ from .sgu_cell import DSGU, DSGUCell, SGU, SGUCell from .sgrn_cell import SGRN, SGRNCell from .star_cell import STAR, STARCell +from .taugru_cell import tauGRU, tauGRUCell from .ugrnn_cell import UGRNN, UGRNNCell from .unicornn_cell import UnICORNN, UnICORNNCell from .wmclstm_cell import WMCLSTM, WMCLSTMCell @@ -99,6 +100,8 @@ "SGRNCell", "STAR", "STARCell", + "tauGRU", + "tauGRUCell", "UGRNN", "UGRNNCell", "UnICORNN", diff --git a/torchrecurrent/cells/taugru_cell.py b/torchrecurrent/cells/taugru_cell.py new file mode 100644 index 0000000..cd3e2c7 --- /dev/null +++ b/torchrecurrent/cells/taugru_cell.py @@ -0,0 +1,314 @@ +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +from torch import Tensor + +from ..base import ( + SingleStateCellBase, + SingleStateRecurrentLayerBase, + resolve_activation, + resolve_init_name, +) + + +class tauGRU(SingleStateRecurrentLayerBase): + r"""Multi-layer tau-GRU neural network. + + [`arXiv `_] + + Each layer consists of a :class:`tauGRUCell`, which updates the hidden state + with a weighted time-delay feedback term: + + .. math:: + \begin{aligned} + u_n &= \tanh(W_1 h_n + U_1 x_n), \\ + z_n &= \tanh(W_2 h_{n-d} + U_2 x_n), \\ + g_n &= \sigma(W_3 h_n + U_3 x_n), \\ + a_n &= \sigma(W_4 h_n + U_4 x_n), \\ + h_{n+1} &= (1 - g_n) \circ h_n + g_n \circ (u_n + a_n \circ z_n), + \end{aligned} + + where :math:`d` is the integer delay in recurrent steps. Delayed states + before the beginning of the sequence are initialized to zero. + + Args: + input_size: The number of expected features in the input `x`. + hidden_size: The number of features in the hidden state `h`. + num_layers: Number of recurrent layers. Default: 1 + dropout: If non-zero, introduces a `Dropout` layer on the outputs of + each layer except the last. Default: 0 + batch_first: If ``True``, input and output tensors are provided as + `(batch, seq, feature)` instead of `(seq, batch, feature)`. + Default: False + delay: Integer delay :math:`d` in recurrent steps. Default: 1 + bias: If ``False``, the layer does not use input-side biases. + Default: True + recurrent_bias: If ``False``, the layer does not use recurrent biases. + Default: True + nonlinearity: Nonlinearity for :math:`u_n` and :math:`z_n`. + Default: :func:`torch.tanh` + gate_nonlinearity: Activation for :math:`g_n` and :math:`a_n`. + Default: :func:`torch.sigmoid` + kernel_init: Initializer for `U_i`. Default: + :func:`torch.nn.init.xavier_uniform_` + recurrent_kernel_init: Initializer for `W_i`. Default: + :func:`torch.nn.init.xavier_uniform_` + bias_init: Initializer for input-side biases. Default: + :func:`torch.nn.init.zeros_` + recurrent_bias_init: Initializer for recurrent biases. Default: + :func:`torch.nn.init.zeros_` + device: The desired device of parameters. + dtype: The desired floating point type of parameters. + + Inputs: input, h_0 + - **input**: tensor of shape :math:`(L, H_{in})` for unbatched input, + :math:`(L, N, H_{in})` when ``batch_first=False`` or + :math:`(N, L, H_{in})` when ``batch_first=True``. + - **h_0**: tensor of shape :math:`(\text{num_layers}, H_{out})` for + unbatched input or :math:`(\text{num_layers}, N, H_{out})`. + Defaults to zeros if not provided. + + Outputs: output, h_n + - **output**: tensor containing the output features from the last layer, + for each timestep. + - **h_n**: tensor containing the final hidden state for each layer. + + Attributes: + cells.{k}.weight_ih : input-hidden weights of shape + `(4*hidden_size, input_size)` for `k = 0`, otherwise + `(4*hidden_size, hidden_size)`. + cells.{k}.weight_hh : hidden-hidden weights of shape + `(4*hidden_size, hidden_size)`. + cells.{k}.bias_ih : input-hidden biases of shape `(4*hidden_size)`. + Only present when ``bias=True``. + cells.{k}.bias_hh : hidden-hidden biases of shape `(4*hidden_size)`. + Only present when ``recurrent_bias=True``. + + .. seealso:: + :class:`tauGRUCell` + """ + + __constants__ = ["delay"] + + delay: int + + def __init__( + self, + input_size: int, + hidden_size: int, + num_layers: int = 1, + dropout: float = 0.0, + batch_first: bool = False, + delay: int = 1, + **kwargs, + ): + if delay < 1: + raise ValueError("delay must be a positive integer.") + super(tauGRU, self).__init__( + input_size, hidden_size, num_layers, dropout, batch_first + ) + self.delay = int(delay) + self.initialize_cells(tauGRUCell, **kwargs) + + def extra_repr(self) -> str: + parts = [super().extra_repr()] + if self.delay != 1: + parts.append(f"delay={self.delay}") + return ", ".join(parts) + + def forward(self, inp: Tensor, state: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]: + if self.batch_first: + inp = inp.transpose(0, 1) + + seq_len, batch_size, _ = inp.size() + + if state is None: + state = torch.zeros( + self.num_layers, + batch_size, + self.hidden_size, + dtype=inp.dtype, + device=inp.device, + ) + + histories = torch.jit.annotate(List[Tensor], []) + for _ in range(self.delay): + histories.append(torch.zeros_like(state)) + + outputs = torch.jit.annotate(List[Tensor], []) + + for t in range(seq_len): + x = inp[t] + new_states = torch.jit.annotate(List[Tensor], []) + delayed_state = histories[0] + + for layer_idx, cell in enumerate(self.cells): + h_prev = state[layer_idx] + h_delay = delayed_state[layer_idx] + h_new = cell(x, h_prev, h_delay) + new_states.append(h_new) + x = h_new + if self.dropout_layer is not None and layer_idx < self.num_layers - 1: + x = self.dropout_layer(x) + + histories = histories[1:] + [state] + state = torch.stack(new_states, dim=0) + outputs.append(x) + + out = torch.stack(outputs, dim=0) + if self.batch_first: + out = out.transpose(0, 1) + return out, state + + +class tauGRUCell(SingleStateCellBase): + r"""A tau-GRU cell with weighted time-delay feedback. + + [`arXiv `_] + + .. math:: + + \begin{aligned} + \mathbf{u}_n &= \phi(W_1 \mathbf{h}_n + U_1 \mathbf{x}_n), \\ + \mathbf{z}_n &= \phi(W_2 \mathbf{h}_{n-d} + U_2 \mathbf{x}_n), \\ + \mathbf{g}_n &= \sigma(W_3 \mathbf{h}_n + U_3 \mathbf{x}_n), \\ + \mathbf{a}_n &= \sigma(W_4 \mathbf{h}_n + U_4 \mathbf{x}_n), \\ + \mathbf{h}_{n+1} &= (1 - \mathbf{g}_n) \circ \mathbf{h}_n + + \mathbf{g}_n \circ + (\mathbf{u}_n + \mathbf{a}_n \circ \mathbf{z}_n). + \end{aligned} + + Args: + input_size: The number of expected features in the input ``x``. + hidden_size: The number of features in the hidden state ``h``. + bias: If ``False``, disables input-side biases. Default: ``True``. + recurrent_bias: If ``False``, disables recurrent biases. Default: ``True``. + nonlinearity: Nonlinearity for the instantaneous and delayed candidates. + Default: :func:`torch.tanh`. + gate_nonlinearity: Activation for the update and feedback gates. + Default: :func:`torch.sigmoid`. + kernel_init: Initializer for ``weight_ih``. Default: + :func:`torch.nn.init.xavier_uniform_`. + recurrent_kernel_init: Initializer for ``weight_hh``. Default: + :func:`torch.nn.init.xavier_uniform_`. + bias_init: Initializer for input-side biases when ``bias=True``. + Default: :func:`torch.nn.init.zeros_`. + recurrent_bias_init: Initializer for recurrent biases when + ``recurrent_bias=True``. Default: :func:`torch.nn.init.zeros_`. + device: The desired device of parameters. + dtype: The desired floating point type of parameters. + + Inputs: input, hidden, delayed_hidden + - **input** of shape ``(batch, input_size)`` or ``(input_size,)``: + tensor containing input features. + - **hidden** of shape ``(batch, hidden_size)`` or ``(hidden_size,)``: + tensor containing the current hidden state. Defaults to zero. + - **delayed_hidden** of shape ``(batch, hidden_size)`` or + ``(hidden_size,)``: tensor containing the delayed hidden state. + Defaults to zero. + + Outputs: h_1 + - **h_1** of shape ``(batch, hidden_size)`` or ``(hidden_size,)``: + tensor containing the next hidden state. + + Variables: + weight_ih: input-hidden weights, of shape ``(4*hidden_size, input_size)`` + weight_hh: hidden-hidden weights, of shape ``(4*hidden_size, hidden_size)`` + bias_ih: input biases, of shape ``(4*hidden_size,)`` if ``bias=True`` + bias_hh: hidden biases, of shape ``(4*hidden_size,)`` if + ``recurrent_bias=True`` + """ + + __constants__ = ["input_size", "hidden_size", "bias", "recurrent_bias"] + + weight_ih: Tensor + weight_hh: Tensor + bias_ih: Tensor + bias_hh: Tensor + + def __init__( + self, + input_size: int, + hidden_size: int, + bias: bool = True, + recurrent_bias: bool = True, + nonlinearity="tanh", + gate_nonlinearity="sigmoid", + kernel_init=nn.init.xavier_uniform_, + recurrent_kernel_init=nn.init.xavier_uniform_, + bias_init=nn.init.zeros_, + recurrent_bias_init=nn.init.zeros_, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + super().__init__( + input_size=input_size, + hidden_size=hidden_size, + bias=bias, + recurrent_bias=recurrent_bias, + device=device, + dtype=dtype, + ) + self.act = resolve_activation(nonlinearity) + self.gate_act = resolve_activation(gate_nonlinearity) + self.init_cfg["kernel"] = resolve_init_name(kernel_init, self.init_cfg["kernel"]) + self.init_cfg["recurrent_kernel"] = resolve_init_name( + recurrent_kernel_init, self.init_cfg["recurrent_kernel"] + ) + self.init_cfg["bias"] = resolve_init_name(bias_init, self.init_cfg["bias"]) + self.init_cfg["recurrent_bias"] = resolve_init_name( + recurrent_bias_init, self.init_cfg["recurrent_bias"] + ) + + self._default_register_tensors(ih_mult=4, hh_mult=4) + self.reset_parameters() + self._cleanup_non_scriptable() + + def forward( + self, + inp: Tensor, + state: Optional[Tensor] = None, + delayed_state: Optional[Tensor] = None, + ) -> Tensor: + self._validate_input(inp) + b_inp, is_batched = self._as_batched(inp) + + if state is None: + b_state = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype) + else: + b_state = state.unsqueeze(0) if (not is_batched and state.dim() == 1) else state + + if delayed_state is None: + b_delayed = self._zeros_state(b_inp.size(0), b_inp.device, b_inp.dtype) + else: + b_delayed = ( + delayed_state.unsqueeze(0) + if (not is_batched and delayed_state.dim() == 1) + else delayed_state + ) + + weight_ih_u, weight_ih_z, weight_ih_g, weight_ih_a = self.weight_ih.chunk(4, 0) + weight_hh_u, weight_hh_z, weight_hh_g, weight_hh_a = self.weight_hh.chunk(4, 0) + bias_ih_u, bias_ih_z, bias_ih_g, bias_ih_a = self.bias_ih.chunk(4, 0) + bias_hh_u, bias_hh_z, bias_hh_g, bias_hh_a = self.bias_hh.chunk(4, 0) + + u = self.act( + b_inp @ weight_ih_u.t() + bias_ih_u + b_state @ weight_hh_u.t() + bias_hh_u + ) + z = self.act( + b_inp @ weight_ih_z.t() + bias_ih_z + b_delayed @ weight_hh_z.t() + bias_hh_z + ) + g = self.gate_act( + b_inp @ weight_ih_g.t() + bias_ih_g + b_state @ weight_hh_g.t() + bias_hh_g + ) + a = self.gate_act( + b_inp @ weight_ih_a.t() + bias_ih_a + b_state @ weight_hh_a.t() + bias_hh_a + ) + + new_state = (1.0 - g) * b_state + g * (u + a * z) + + if not is_batched: + new_state = new_state.squeeze(0) + + return new_state From 1bc10f338bba5317612346619c1f8f8d7a4dded0 Mon Sep 17 00:00:00 2001 From: MartinuzziFrancesco Date: Thu, 25 Jun 2026 10:52:20 +0200 Subject: [PATCH 2/2] chore: up version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3db5781..2a5687f 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "torchrecurrent" -version = "0.2.2" +version = "0.2.3" description = "A package for recurrent neural networks in PyTorch" readme = "README.md" authors = [