[host_callback] Remove most of the jax.experimental.host_callback module

These APIs have been deprecated since March 2024 and they are subsumed by the new JAX external callbacks. See #20385 for a discussion. PiperOrigin-RevId: 682659677
jax-ml · Oct 6, 2024 · 1d31e2f · 1d31e2f
1 parent 81a31f6
commit 1d31e2f
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 4,077 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -16,6 +16,9 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * {func}`jax.numpy.isscalar` now returns True for any array-like object with
     zero dimensions. Previously it only returned True for zero-dimensional
     array-like objects with a weak dtype.
+  * `jax.experimental.host_callback` has been deprecated since March 2024, with
+    JAX version 0.4.26. Now we removed it.
+    See {jax-issue}`#20385` for a discussion of alternatives.
 
 ## jax 0.4.34 (October 4, 2023)
 

diff --git a/docs/jax.experimental.host_callback.rst b/docs/jax.experimental.host_callback.rst
diff --git a/docs/jax.experimental.rst b/docs/jax.experimental.rst
@@ -16,7 +16,6 @@ Experimental Modules
 
     jax.experimental.array_api
     jax.experimental.checkify
-    jax.experimental.host_callback
     jax.experimental.pjit
     jax.experimental.sparse
     jax.experimental.jet