Remove remaining implementations of jax.experimental.host_callback.call.

The host_callback module has been deprecated since March 2024, and we are now removing the implementation. We keep the functions so that we can give a nicer error message than AttributeError, and because removing those now break internal pytype checking. We will remove those in the near future.

See #20385.

PiperOrigin-RevId: 682837548

jax/experimental/host_callback.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Backwards compatibility shim for the deprecated host_callback APIs.
+"""Backwards compatibility shims for the deprecated host_callback APIs.
 
 .. warning::
   The host_callback APIs are deprecated as of March 20, 2024.
@@ -23,64 +23,12 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
-import logging
-import warnings
+def call(*_, **__):
+  raise NotImplementedError(
+      "jax.experimental.host_callback has been deprecated since March 2024 and "
+      "is now no longer supported. "
+      "See https://github.com/jax-ml/jax/issues/20385"
+  )
 
-import jax
-from jax.experimental import io_callback
 
-
-logger = logging.getLogger(__name__)
-
-
-# We keep a shim for host_callback.call because it is still used in a few
-# places in google.
-def call(callback_func: Callable,
-         arg,
-         *,
-         result_shape=None,
-         call_with_device=False,
-         device_index=0,
-         callback_flavor=None):
-  """Make a call to the host, and expect a result.
-
-  .. warning::
-    The host_callback APIs are deprecated as of March 20, 2024.
-    The functionality is subsumed by the
-    `new JAX external callbacks <https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html>`_
-    See https://github.com/jax-ml/jax/issues/20385.
-  """
-  warnings.warn("""The host_callback APIs are deprecated as of March 20, 2024.
-    The functionality is subsumed by the
-    new JAX external callbacks (https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html).
-    See https://github.com/jax-ml/jax/issues/20385
-  """, DeprecationWarning, stacklevel=2)
-  if callback_flavor is not None:
-    raise NotImplementedError(
-        "host_callback.call is only supported with the IO_CALLBACK flavor.")
-  if call_with_device:
-    raise NotImplementedError(
-        "host_callback.call is only supported with the call_with_device=False.")
-  callback_device = jax.local_devices()[device_index]
-  sharding = jax.sharding.SingleDeviceSharding(callback_device)
-  return io_callback(callback_func, result_shape, arg,
-                     sharding=sharding,
-                     ordered=True)
-
-import typing
-if typing.TYPE_CHECKING:
-  def id_tap(tap_func,
-            arg,
-            *,
-            result=None,
-            tap_with_device=False,
-            device_index=0,
-            callback_flavor=None,
-            **kwargs):
-    raise NotImplementedError(
-        "host_callback.id_tap is no longer supported. "
-        "See https://github.com/jax-ml/jax/issues/20385"
-    )
-
-del typing
+id_tap = call

tests/BUILD

@@ -330,7 +330,6 @@ jax_multiplatform_test(
     name = "infeed_test",
     srcs = ["infeed_test.py"],
     deps = [
-        "//jax:experimental_host_callback",
     ],
 )
 
@@ -1142,16 +1141,6 @@ jax_multiplatform_test(
     deps = ["//jax:ode"],
 )
 
-jax_multiplatform_test(
-    name = "host_callback_test",
-    srcs = ["host_callback_test.py"],
-    main = "host_callback_test.py",
-    deps = [
-        "//jax:experimental",
-        "//jax:ode",
-    ],
-)
-
 jax_multiplatform_test(
     name = "key_reuse_test",
     srcs = ["key_reuse_test.py"],