feat: add taugrucell

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

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

docs/generated/torchrecurrent.tauGRU.rst

@@ -0,0 +1,8 @@
+torchrecurrent.tauGRU
+=====================
+
+.. currentmodule:: torchrecurrent
+
+.. autoclass:: tauGRU
+   :members:
+   :show-inheritance:

docs/generated/torchrecurrent.tauGRUCell.rst

@@ -0,0 +1,8 @@
+torchrecurrent.tauGRUCell
+=========================
+
+.. currentmodule:: torchrecurrent
+
+.. autoclass:: tauGRUCell
+   :members:
+   :show-inheritance:

docs/models.rst

@@ -110,6 +110,9 @@ references and official implementations where available.
    * - :doc:`STAR <generated/torchrecurrent.STAR>`
      - `TPAMI 2022 <https://arxiv.org/abs/1911.11033>`__
      - `0zgur0/STAckable-Recurrent-network <https://github.com/0zgur0/STAckable-Recurrent-network>`__
+   * - :doc:`tauGRU <generated/torchrecurrent.tauGRU>`
+     - `AISTATS 2025 <https://arxiv.org/abs/2212.00228>`__
+     - –
    * - :doc:`UGRNN <generated/torchrecurrent.UGRNN>`
      - `ICLR 2017 <https://arxiv.org/abs/1611.09913>`__
      - –

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 = [

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."""

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),

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",

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",