Repository: google/jax
Pull request #21321 (2024-05-20 22:30:57)

Refactor & test internal deprecation APIs
Refactor & test internal deprecation APIs

The names and APIs were previously too similar and therefore somewhat confusing; this will be more clear I think.



Patch:
diff --git a/jax/__init__.py b/jax/__init__.py
index f274c0e37..f3ee0edd7 100644
--- a/jax/__init__.py
+++ b/jax/__init__.py
@@ -179,7 +179,7 @@ import jax.experimental.compilation_cache.compilation_cache as _ccache
 del _ccache
 
 from jax._src.deprecations import register as _register_deprecation
-_register_deprecation("jax.experimental", "maps-module")
+_register_deprecation("jax-experimental-maps-module")
 del _register_deprecation
 
 _deprecations = {
diff --git a/jax/_src/deprecations.py b/jax/_src/deprecations.py
index cd44f1c26..0bbae2d14 100644
--- a/jax/_src/deprecations.py
+++ b/jax/_src/deprecations.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
 from types import ModuleType
 import warnings
 
@@ -56,7 +57,7 @@ def deprecation_getattr(module, deprecations):
   return getattr
 
 
-def accelerate_module_deprecation(module: ModuleType, name: str) -> None:
+def accelerate_getattr_deprecation(module: ModuleType, name: str) -> None:
   """Accelerate the deprecation of a module-level attribute.
 
   Raises an AttributeError instead of a DeprecationWarning upon attribute access.
@@ -68,24 +69,35 @@ def accelerate_module_deprecation(module: ModuleType, name: str) -> None:
 # The following mechanism is a separate one, for registering and
 # accelerating deprecations that are not imports (for example, deprecations
 # of a function argument).
-# Maps a pair of strings to a boolean specifying whether the deprecation
-# is accelerated. The intent is that non-accelerated deprecations will warn,
-# and accelerated deprecations will error.
-_registered_deprecations: dict[tuple[str, str], bool] = {}
+# Maps a globally unique string ID to a DeprecationState, which tracks whether
+# the deprecation is accelerated.
+# The intent is that non-accelerated deprecations will warn, and accelerated
+# deprecations will error.
 
+@dataclass
+class DeprecationState:
+  accelerated: bool = False
 
-def register(module: str, key: str) -> None:
-  _registered_deprecations[module, key] = False
+_registered_deprecations: dict[str, DeprecationState] = {}
 
 
-def unregister(module: str, key: str) -> None:
-  _registered_deprecations.pop((module, key))
+def register(deprecation_id: str) -> None:
+  _registered_deprecations[deprecation_id] = DeprecationState()
 
 
-def accelerate(module: str, key: str) -> None:
-  assert (module, key) in _registered_deprecations
-  _registered_deprecations[module, key] = True
+def unregister(deprecation_id: str) -> None:
+  if deprecation_id not in _registered_deprecations:
+    raise ValueError(f"{deprecation_id=!r} not registered.")
+  _registered_deprecations.pop(deprecation_id)
 
 
-def is_accelerated(module: str, key: str) -> bool:
-  return _registered_deprecations[module, key]
+def accelerate(deprecation_id: str) -> None:
+  if deprecation_id not in _registered_deprecations:
+    raise ValueError(f"{deprecation_id=!r} not registered.")
+  _registered_deprecations[deprecation_id].accelerated = True
+
+
+def is_accelerated(deprecation_id: str) -> bool:
+  if deprecation_id not in _registered_deprecations:
+    raise ValueError(f"{deprecation_id=!r} not registered.")
+  return _registered_deprecations[deprecation_id].accelerated
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 6aa5bbdb9..f9029e648 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2704,7 +2704,7 @@ available in the JAX FAQ at :ref:`faq-data-placement` (full FAQ at
 https://jax.readthedocs.io/en/latest/faq.html).
 """
 
-deprecations.register(__name__, "array-none")
+deprecations.register("jax-numpy-array-none")
 
 @util.implements(np.array, lax_description=_ARRAY_DOC)
 def array(object: Any, dtype: DTypeLike | None = None, copy: bool = True,
@@ -2762,7 +2762,7 @@ def array(object: Any, dtype: DTypeLike | None = None, copy: bool = True,
   leaves = tree_leaves(object, is_leaf=lambda x: x is None)
   if any(leaf is None for leaf in leaves):
     # Added Nov 16 2023
-    if deprecations.is_accelerated(__name__, "array-none"):
+    if deprecations.is_accelerated("jax-numpy-array-none"):
       raise TypeError("None is not a valid value for jnp.array")
     warnings.warn(
       "None encountered in jnp.array(); this is currently treated as NaN. "
diff --git a/jax/experimental/host_callback.py b/jax/experimental/host_callback.py
index 3e591ec6d..55b742a30 100644
--- a/jax/experimental/host_callback.py
+++ b/jax/experimental/host_callback.py
@@ -2014,9 +2014,6 @@ if typing.TYPE_CHECKING:
   stop_outfeed_receiver = _deprecated_stop_outfeed_receiver
 else:
   from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
-  from jax._src.deprecations import register
-  for deprecated in _deprecations.keys():
-    register(__name__, deprecated)
   __getattr__ = _deprecation_getattr(__name__, _deprecations)
   del _deprecation_getattr
 del typing
diff --git a/jax/experimental/maps.py b/jax/experimental/maps.py
index ab2d90e76..e378a6fc2 100644
--- a/jax/experimental/maps.py
+++ b/jax/experimental/maps.py
@@ -41,7 +41,7 @@ _msg = (
     " jax.experimental.shard_map nor jax.vmap are suitable for your use case."
 )
 
-if deprecations.is_accelerated("jax.experimental", "maps-module"):
+if deprecations.is_accelerated("jax-experimental-maps-module"):
   raise ImportError(_msg)
 else:
   warnings.warn(_msg, DeprecationWarning, stacklevel=2)


Test patch:
diff --git a/tests/deprecation_test.py b/tests/deprecation_test.py
index 02d9c427c..2476a6208 100644
--- a/tests/deprecation_test.py
+++ b/tests/deprecation_test.py
@@ -15,12 +15,13 @@
 import warnings
 
 from absl.testing import absltest
+from jax._src import deprecations
 from jax._src import test_util as jtu
 from jax._src.internal_test_util import deprecation_module as m
 
 class DeprecationTest(absltest.TestCase):
 
-  def testDeprecation(self):
+  def testModuleDeprecation(self):
     with warnings.catch_warnings():
       warnings.simplefilter("error")
       self.assertEqual(m.x, 42)
@@ -35,6 +36,24 @@ class DeprecationTest(absltest.TestCase):
                                 "module .* has no attribute 'w'"):
       _ = m.w
 
+  def testNamedDeprecation(self):
+    some_unique_id = "some-unique-id"
+    try:
+      deprecations.register(some_unique_id)
+      self.assertFalse(deprecations.is_accelerated(some_unique_id))
+      deprecations.accelerate(some_unique_id)
+      self.assertTrue(deprecations.is_accelerated(some_unique_id))
+    finally:
+      deprecations.unregister(some_unique_id)
+
+    msg = f"deprecation_id={some_unique_id!r} not registered"
+    with self.assertRaisesRegex(ValueError, msg):
+      deprecations.accelerate(some_unique_id)
+    with self.assertRaisesRegex(ValueError, msg):
+      deprecations.is_accelerated(some_unique_id)
+    with self.assertRaisesRegex(ValueError, msg):
+      deprecations.unregister(some_unique_id)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())