From c90f1f0c960da87716d5c49e02efc4571e8a0781 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Fri, 27 Oct 2023 09:40:23 +0100
Subject: [PATCH] Deprecated accessing `config` via the `jax.config` submodule

The `config` object it re-exports is available on the top-level `jax` package,
i.e.

    from jax.config import config

can always safely be replaced with

    from jax import config

or just

    import jax
    jax.config
---
 jax/config.py | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/jax/config.py b/jax/config.py
index 947c5abf8230..9435308d157f 100644
--- a/jax/config.py
+++ b/jax/config.py
@@ -12,6 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# TODO(phawkins): fix users of this alias and delete this file.
+from jax._src.config import config as _deprecated_config  # noqa: F401
 
-from jax._src.config import config  # noqa: F401
+# Deprecations
+
+_deprecations = {
+    # Added October 27, 2023
+    "config": (
+        "Accessing jax.config via the jax.config submodule is deprecated.",
+        _deprecated_config),
+}
+
+import typing
+if typing.TYPE_CHECKING:
+  config = _deprecated_config
+else:
+  from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
+  __getattr__ = _deprecation_getattr(__name__, _deprecations)
+  del _deprecation_getattr
+del typing
+del _deprecated_config