Repository: google/flax

Issue:
Feature request: Add minimal gated unit (MGU)
Feature request: Add [minimal gated unit (MGU)](https://en.wikipedia.org/wiki/Gated_recurrent_unit#Minimal_gated_unit) to [flax.linen](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html#rnn-primitives). I can open a PR if desired.


Discussion:

Resolving patch:
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -111,6 +111,7 @@
     Bidirectional as Bidirectional,
     ConvLSTMCell as ConvLSTMCell,
     GRUCell as GRUCell,
+    MGUCell as MGUCell,
     LSTMCell as LSTMCell,
     OptimizedLSTMCell as OptimizedLSTMCell,
     RNNCellBase as RNNCellBase,
diff --git a/flax/linen/recurrent.py b/flax/linen/recurrent.py
--- a/flax/linen/recurrent.py
+++ b/flax/linen/recurrent.py
@@ -462,6 +462,111 @@ def num_feature_axes(self) -> int:
     return 1
 
 
+class MGUCell(RNNCellBase):
+  r"""MGU cell (https://arxiv.org/pdf/1603.09420.pdf).
+
+  The mathematical definition of the cell is as follows
+
+  .. math::
+
+      \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 \\
+      \end{array}
+
+  where x is the input and h, is the output of the previous time step.
+
+  Attributes:
+    features: number of output features.
+    gate_fn: activation function used for gates (default: sigmoid).
+    activation_fn: activation function used for output and memory update
+      (default: tanh).
+    kernel_init: initializer function for the kernels that transform
+      the input (default: lecun_normal).
+    recurrent_kernel_init: initializer function for the kernels that transform
+      the hidden state (default: initializers.orthogonal()).
+    forget_bias_init: initializer for the bias parameters of the forget gate.
+      The default is set to initializers.ones_init() because this prevents
+      vanishing gradients. See https://proceedings.mlr.press/v37/jozefowicz15.pdf,
+      section 2.2 for more details.
+    activation_bias_init: initializer for the bias parameters of the activation
+      output (default: initializers.zeros_init()).
+    dtype: the dtype of the computation (default: None).
+    param_dtype: the dtype passed to parameter initializers (default: float32).
+  """
+  features: int
+  gate_fn: Callable[..., Any] = sigmoid
+  activation_fn: Callable[..., Any] = tanh
+  kernel_init: initializers.Initializer = default_kernel_init
+  recurrent_kernel_init: initializers.Initializer = initializers.orthogonal()
+  forget_bias_init: initializers.Initializer = initializers.ones_init()
+  activation_bias_init: initializers.Initializer = initializers.zeros_init()
+  dtype: Optional[Dtype] = None
+  param_dtype: Dtype = jnp.float32
+  carry_init: initializers.Initializer = initializers.zeros_init()
+
+  @compact
+  def __call__(self, carry, inputs):
+    """Minimal gated unit (MGU) cell.
+
+    Args:
+      carry: the hidden state of the MGU cell,
+        initialized using `MGUCell.initialize_carry`.
+      inputs: an ndarray with the input for the current time step.
+        All dimensions except the final are considered batch dimensions.
+
+    Returns:
+      A tuple with the new carry and the output.
+    """
+    h = carry
+    hidden_features = h.shape[-1]
+    # input and recurrent layers are summed so only one needs a bias.
+    dense_h = partial(
+        Dense,
+        features=hidden_features,
+        use_bias=False,
+        dtype=self.dtype,
+        param_dtype=self.param_dtype,
+        kernel_init=self.recurrent_kernel_init,
+        bias_init=self.activation_bias_init,
+    )
+    dense_i = partial(
+        Dense,
+        features=hidden_features,
+        use_bias=True,
+        dtype=self.dtype,
+        param_dtype=self.param_dtype,
+        kernel_init=self.kernel_init,
+    )
+    f = self.gate_fn(dense_i(name='if', bias_init=self.forget_bias_init)(inputs) + dense_h(name='hf')(h))
+    # add bias because the linear transformations aren't directly summed.
+    n = self.activation_fn(
+        dense_i(name='in', bias_init=self.activation_bias_init)(inputs) + f * dense_h(name='hn', use_bias=True)(h)
+    )
+    new_h = (1.0 - f) * n + f * h
+    return new_h, new_h
+
+  @nowrap
+  def initialize_carry(self, rng: PRNGKey, input_shape: Tuple[int, ...]):
+    """Initialize the RNN cell carry.
+
+    Args:
+      rng: random number generator passed to the init_fn.
+      input_shape: a tuple providing the shape of the input to the cell.
+
+    Returns:
+      An initialized carry for the given RNN cell.
+    """
+    batch_dims = input_shape[:-1]
+    mem_shape = batch_dims + (self.features,)
+    return self.carry_init(rng, mem_shape, self.param_dtype)
+
+  @property
+  def num_feature_axes(self) -> int:
+    return 1
+
+
 class ConvLSTMCell(RNNCellBase):
   r"""A convolutional LSTM cell.
 


Test patch:
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -747,7 +747,7 @@ def __call__(self, x):
 
 
 # TODO(flax-dev): add integration tests for RNN cells
-class RecurrentTest(absltest.TestCase):
+class RecurrentTest(parameterized.TestCase):
 
   def test_lstm(self):
     lstm = nn.LSTMCell(features=4)
@@ -776,37 +776,56 @@ def test_lstm(self):
         },
     )
 
-  def test_gru(self):
-    gru = nn.GRUCell(features=4)
+  @parameterized.parameters(
+      {'module_cls': nn.GRUCell,
+       'expected_param_shapes': {
+         'ir': {'kernel': (3, 4), 'bias': (4,)},
+         'iz': {'kernel': (3, 4), 'bias': (4,)},
+         'in': {'kernel': (3, 4), 'bias': (4,)},
+         'hr': {'kernel': (4, 4)},
+         'hz': {'kernel': (4, 4)},
+         'hn': {'kernel': (4, 4), 'bias': (4,)},
+        }},
+      {'module_cls': nn.MGUCell,
+       'expected_param_shapes': {
+         'if': {'kernel': (3, 4), 'bias': (4,)},
+         'in': {'kernel': (3, 4), 'bias': (4,)},
+         'hf': {'kernel': (4, 4)},
+         'hn': {'kernel': (4, 4), 'bias': (4,)},
+        }}
+  )
+  def test_gated_units(self, module_cls, expected_param_shapes):
+    module = module_cls(features=4)
     rng = random.key(0)
     key1, key2 = random.split(rng)
     x = random.normal(key1, (2, 3))
-    carry0 = gru.initialize_carry(rng, x.shape)
+    carry0 = module.initialize_carry(rng, x.shape)
     self.assertEqual(carry0.shape, (2, 4))
-    (carry, y), initial_params = gru.init_with_output(key2, carry0, x)
+    (carry, y), initial_params = module.init_with_output(key2, carry0, x)
     self.assertEqual(carry.shape, (2, 4))
     np.testing.assert_allclose(y, carry)
     param_shapes = jax.tree_util.tree_map(np.shape, initial_params['params'])
     self.assertEqual(
         param_shapes,
-        {
-            'ir': {'kernel': (3, 4), 'bias': (4,)},
-            'iz': {'kernel': (3, 4), 'bias': (4,)},
-            'in': {'kernel': (3, 4), 'bias': (4,)},
-            'hr': {'kernel': (4, 4)},
-            'hz': {'kernel': (4, 4)},
-            'hn': {'kernel': (4, 4), 'bias': (4,)},
-        },
+        expected_param_shapes,
     )
+    if module_cls == nn.MGUCell:
+      self.assertTrue((initial_params['params']['if']['bias'] == jnp.ones((4,))).all())
+      self.assertTrue((initial_params['params']['in']['bias'] == jnp.zeros((4,))).all())
+      self.assertTrue((initial_params['params']['hn']['bias'] == jnp.zeros((4,))).all())
 
-  def test_complex_input_gru(self):
-    gru = nn.GRUCell(features=4)
+  @parameterized.parameters(
+      {'module_cls': nn.GRUCell},
+      {'module_cls': nn.MGUCell}
+  )
+  def test_complex_input_gated_units(self, module_cls):
+    module_instance = module_cls(features=4)
     rng = random.key(0)
     key1, key2 = random.split(rng)
     x = random.normal(key1, (2, 3), dtype=jnp.complex64)
-    carry0 = gru.initialize_carry(rng, x.shape)
+    carry0 = module_instance.initialize_carry(rng, x.shape)
     self.assertEqual(carry0.shape, (2, 4))
-    (carry, y), _ = gru.init_with_output(key2, carry0, x)
+    (carry, y), _ = module_instance.init_with_output(key2, carry0, x)
     self.assertEqual(carry.dtype, jnp.complex64)
     self.assertEqual(y.dtype, jnp.complex64)