added MGU reset_gate test

docs/api_reference/flax.linen/layers.rst

@@ -183,6 +183,7 @@ BatchApply
   OptimizedLSTMCell
   SimpleCell
   GRUCell
+  MGUCell
   RNN
   Bidirectional
   BatchApply

tests/linen/linen_test.py

@@ -1142,6 +1142,30 @@ def test_optimized_lstm_cell_matches_regular(self):
     np.testing.assert_allclose(y, y_opt, rtol=1e-6)
     check_eq(lstm_params, lstm_opt_params)
 
+  def test_mgu_reset_gate(self):
+    module = nn.MGUCell(features=4, reset_gate=False)
+    rng = random.key(0)
+    rng, key1, key2 = random.split(rng, 3)
+    x = random.normal(key1, (2, 3))
+    carry0 = module.initialize_carry(rng, x.shape)
+    (carry, y), v = module.init_with_output(key2, carry0, x)
+
+    self.assertIn('kernel', v['params']['hn'])
+    self.assertNotIn('bias', v['params']['hn'])
+
+    f = jax.nn.sigmoid(
+      jnp.dot(x, v['params']['if']['kernel'])
+      + v['params']['if']['bias'].reshape(1, -1)
+      + jnp.dot(carry0, v['params']['hf']['kernel'])
+    )
+    n = jax.nn.tanh(
+      jnp.dot(x, v['params']['in']['kernel'])
+      + v['params']['in']['bias'].reshape(1, -1)
+      + jnp.dot(carry0, v['params']['hn']['kernel'])
+    )
+    expected_out = (1 - f) * n + f * carry0
+    np.testing.assert_allclose(y, expected_out)
+
 
 class IdsTest(absltest.TestCase):
   def test_hashable(self):