Fix MGUCell hidden state update

flax/linen/recurrent.py

@@ -624,7 +624,7 @@ class MGUCell(RNNCellBase):
       \begin{array}{ll}
       f = \sigma(W_{if} x + b_{if} + W_{hf} h) \\
       n = \tanh(W_{in} x + b_{in} + f * (W_{hn} h + b_{hn})) \\
-      h' = (1 - f) * n + f * h \\
+      h' = (1 - f) * h + f * n \\
       \end{array}
 
   where x is the input and h is the output of the previous time step.
@@ -636,7 +636,7 @@ class MGUCell(RNNCellBase):
       \begin{array}{ll}
       f = \sigma(W_{if} x + b_{if} + W_{hf} h) \\
       n = \tanh(W_{in} x + b_{in} + W_{hn} h) \\
-      h' = (1 - f) * n + f * h \\
+      h' = (1 - f) * h + f * n \\
       \end{array}
 
   Example usage::
@@ -726,7 +726,7 @@ def __call__(self, carry, inputs):
     n = self.activation_fn(
       dense_i(name="in", bias_init=self.activation_bias_init)(inputs) + x
     )
-    new_h = (1.0 - f) * n + f * h
+    new_h = (1.0 - f) * h + f * n
     return new_h, new_h
 
   @nowrap

tests/linen/linen_test.py

@@ -1243,29 +1243,41 @@ 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)
+  @parameterized.parameters(
+    {'reset_gate': True},
+    {'reset_gate': False},
+  )
+  def test_mgu_cell_matches_manual_recurrence(self, reset_gate):
+    module = nn.MGUCell(features=4, reset_gate=reset_gate)
     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'])
+    if reset_gate:
+      self.assertIn('bias', v['params']['hn'])
+    else:
+      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'])
     )
+    recurrent_update = jnp.dot(carry0, v['params']['hn']['kernel'])
+    if reset_gate:
+      recurrent_update += v['params']['hn']['bias'].reshape(1, -1)
+      recurrent_update *= f
     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'])
+      + recurrent_update
     )
-    expected_out = (1 - f) * n + f * carry0
+    expected_out = (1 - f) * carry0 + f * n
     np.testing.assert_allclose(y, expected_out)
+    np.testing.assert_allclose(carry, expected_out)
 
 
 class IdsTest(absltest.TestCase):