From 251dfcad3c0a344e1eda0532862072a8a7932416 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <jakevdp@google.com>
Date: Fri, 28 Jun 2024 10:47:37 -0700
Subject: [PATCH] Deprecate jax.interpreters xb, xc, xe abbreviations.

Instead, import directly as jax.lib.xla_bridge, jax.lib.xla_client, jax.lib.xla_extension.
---
 jax/interpreters/xla.py | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 79a736ae3091..bbd5b65d5d3e 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -23,14 +23,27 @@
   apply_primitive as apply_primitive,
 )
 
-from jax._src import xla_bridge as xb
-from jax._src.lib import xla_client as xc
+from jax._src import xla_bridge as _xb
+from jax._src.lib import xla_client as _xc
 
-xe = xc._xla
-Backend = xe.Client
+_xe = _xc._xla
+Backend = _xe.Client
 
 # Deprecations
 _deprecations = {
+    # Added 2024-06-28
+    "xb": (
+        "jax.interpreters.xla.xb is deprecated. Use jax.lib.xla_bridge instead.",
+        _xb
+    ),
+    "xc": (
+        "jax.interpreters.xla.xc is deprecated. Use jax.lib.xla_client instead.",
+        _xc,
+    ),
+    "xe": (
+        "jax.interpreters.xla.xe is deprecated. Use jax.lib.xla_extension instead.",
+        _xe,
+    ),
     # Finalized 2024-05-13; remove after 2024-08-13
     "backend_specific_translations": (
         "jax.interpreters.xla.backend_specific_translations is deprecated. "
@@ -69,6 +82,13 @@
     ),
 }
 
+import typing
 from jax._src.deprecations import deprecation_getattr as _deprecation_getattr
-__getattr__ = _deprecation_getattr(__name__, _deprecations)
+if typing.TYPE_CHECKING:
+  xb = _xb
+  xc = _xc
+  xe = _xe
+else:
+  __getattr__ = _deprecation_getattr(__name__, _deprecations)
 del _deprecation_getattr
+del typing