Repository: google/flax

Issue:
Remove RNNCellBase refactor error messages
During the `RNNCellBase` refactor some tooling was created to produce error messages if users tried to leverage the old API. Now that a few months have passed, it would be a bit safer to assume that most active users have seen that error message and have fixed their code so we can remove this tooling.

https://github.com/google/flax/blob/0347dd2f34b1a581dd664698b7c842ca2f650745/flax/linen/recurrent.py#L60-L66


Discussion:

Resolving patch:
diff --git a/flax/linen/recurrent.py b/flax/linen/recurrent.py
--- a/flax/linen/recurrent.py
+++ b/flax/linen/recurrent.py
@@ -18,7 +18,6 @@
 see: https://flax.readthedocs.io/en/latest/developer_notes/lift.html.
 """
 
-from abc import ABCMeta
 from functools import partial  # pylint: disable=g-importing-member
 from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, TypeVar, Union
 from absl import logging
@@ -51,41 +50,6 @@
 Output = Any
 
 
-class _Never:
-  pass
-
-
-NEVER = _Never()
-
-LEGACY_UPDATE_MESSAGE = (
-    'The RNNCellBase API has changed, '
-    'the error you are experiencing might be caused by this change. Please '
-    'update your code to the new API, for more information on how to do this '
-    'please check out the RNNCellBase migration guide: '
-    'https://flax.readthedocs.io/en/latest/guides/rnncell_upgrade_guide.html'
-)
-
-
-class RNNCellCompatibilityMeta(ABCMeta):
-  """Metaclass for RNNCell compatibility."""
-
-  def __call__(self, *args: Any, **kwds: Any) -> Any:
-    try:
-      return super().__call__(*args, **kwds)
-    except TypeError as e:
-      msg = e.args[0]
-      raise TypeError(f'{msg} \n\n {LEGACY_UPDATE_MESSAGE}') from e
-
-
-def deprecation_method_decorator(f):
-  def wrapper(*args, **kwargs):
-    if len(args) < 1 or not isinstance(args[0], RNNCellBase):
-      raise TypeError(LEGACY_UPDATE_MESSAGE)
-    return f(*args, **kwargs)
-
-  return wrapper
-
-
 class RNNCellBase(Module):
   """RNN cell base class."""
 
@@ -110,7 +74,7 @@ def num_feature_axes(self) -> int:
     raise NotImplementedError
 
 
-class LSTMCell(RNNCellBase, metaclass=RNNCellCompatibilityMeta):
+class LSTMCell(RNNCellBase):
   r"""LSTM cell.
 
   The mathematical definition of the cell is as follows
@@ -193,7 +157,6 @@ def __call__(self, carry, inputs):
     return (new_c, new_h), new_h
 
   @nowrap
-  @deprecation_method_decorator
   def initialize_carry(
       self, rng: PRNGKey, input_shape: Tuple[int, ...]
   ) -> Tuple[Array, Array]:
@@ -244,7 +207,7 @@ def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]:
     return k, b
 
 
-class OptimizedLSTMCell(RNNCellBase, metaclass=RNNCellCompatibilityMeta):
+class OptimizedLSTMCell(RNNCellBase):
   r"""More efficient LSTM Cell that concatenates state components before matmul.
 
   The parameters are compatible with `LSTMCell`. Note that this cell is often
@@ -373,7 +336,6 @@ def _concat_dense(
     return (new_c, new_h), new_h
 
   @nowrap
-  @deprecation_method_decorator
   def initialize_carry(
       self, rng: PRNGKey, input_shape: Tuple[int, ...]
   ) -> Tuple[Array, Array]:
@@ -398,7 +360,7 @@ def num_feature_axes(self) -> int:
     return 1
 
 
-class GRUCell(RNNCellBase, metaclass=RNNCellCompatibilityMeta):
+class GRUCell(RNNCellBase):
   r"""GRU cell.
 
   The mathematical definition of the cell is as follows
@@ -481,7 +443,6 @@ def __call__(self, carry, inputs):
     return new_h, new_h
 
   @nowrap
-  @deprecation_method_decorator
   def initialize_carry(self, rng: PRNGKey, input_shape: Tuple[int, ...]):
     """Initialize the RNN cell carry.
 
@@ -596,7 +557,6 @@ def __call__(self, carry, inputs):
     return (new_c, new_h), new_h
 
   @nowrap
-  @deprecation_method_decorator
   def initialize_carry(self, rng: PRNGKey, input_shape: Tuple[int, ...]):
     """Initialize the RNN cell carry.
 
@@ -723,7 +683,6 @@ class RNN(Module):
   """
 
   cell: RNNCellBase
-  cell_size: Any = NEVER
   time_major: bool = False
   return_carry: bool = False
   reverse: bool = False
@@ -738,14 +697,6 @@ class RNN(Module):
       {'params': False}
   )
 
-  def __post_init__(self) -> None:
-    if self.cell_size is not NEVER:
-      raise TypeError(
-          f'The `cell_size` argument is no longer available`. '
-          + LEGACY_UPDATE_MESSAGE
-      )
-    return super().__post_init__()
-
   def __call__(
       self,
       inputs: jax.Array,


Test patch:
diff --git a/tests/linen/linen_recurrent_test.py b/tests/linen/linen_recurrent_test.py
--- a/tests/linen/linen_recurrent_test.py
+++ b/tests/linen/linen_recurrent_test.py
@@ -15,7 +15,7 @@
 """Recurrent tests."""
 
 
-from absl.testing import absltest, parameterized
+from absl.testing import absltest
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -511,28 +511,5 @@ def test_return_carry(self):
     )
 
 
-class TestRecurrentDeprecation(parameterized.TestCase):
-
-  @parameterized.product(
-      cell_type=[nn.LSTMCell, nn.GRUCell, nn.OptimizedLSTMCell]
-  )
-  def test_constructor(self, cell_type):
-    with self.assertRaisesRegex(TypeError, 'The RNNCellBase API has changed'):
-      cell_type()
-
-  @parameterized.product(
-      cell_type=[nn.LSTMCell, nn.GRUCell, nn.OptimizedLSTMCell]
-  )
-  def test_initialize_carry(self, cell_type):
-    key = jax.random.key(0)
-    with self.assertRaisesRegex(TypeError, 'The RNNCellBase API has changed'):
-      cell_type.initialize_carry(key, (2,), 3)
-
-  def test_rnn(self):
-    cell = nn.LSTMCell(3)
-    with self.assertRaisesRegex(TypeError, 'The RNNCellBase API has changed'):
-      nn.RNN(cell, cell_size=8)
-
-
 if __name__ == '__main__':
   absltest.main()