Register maps module deprecation outside of module

PiperOrigin-RevId: 617194807

Author: Jake VanderPlas

jax/__init__.py

@@ -179,6 +179,7 @@
 # TODO(jakevdp): remove this when jax/config.py is removed.
 from jax._src.deprecations import register as _register_deprecation
 _register_deprecation("jax.config", "config-module")
+_register_deprecation("jax.experimental", "maps-module")
 del _register_deprecation
 
 _deprecations = {

jax/experimental/maps.py

@@ -41,9 +41,7 @@
     " jax.experimental.shard_map nor jax.vmap are suitable for your use case."
 )
 
-deprecations.register("jax.experimental.maps", "maps-module")
-
-if deprecations.is_accelerated("jax.experimental.maps", "maps-module"):
+if deprecations.is_accelerated("jax.experimental", "maps-module"):
   raise ImportError(_msg)
 else:
   warnings.warn(_msg, DeprecationWarning, stacklevel=2)