Repository: jax-ml/jax
Pull request #27841 (2025-04-08 18:47:48)

Finalize deprecation of vectorized argument in callbacks.
Finalize deprecation of vectorized argument in callbacks.

The `vectorized` argument to `pure_callback` and `ffi_call` was deprecated in JAX v0.4.34 (released Oct 4 2024), then added to the CHANGELOG in v0.4.35 (doh! released Oct 22). The JAX compatibility policy requires 3 months of compatible releases before a deprecation is finalized, so it is time to remove this parameter from the public API. The `vmap_method` parameter can be used instead, and the docs for [`pure_callback`](https://docs.jax.dev/en/latest/_autosummary/jax.pure_callback.html) provide more details.

This change has one other (non-obvious!) affect on the user facing APIs. (Note that this change in behavior has also been protected by a deprecation warning since the `vectorized` parameter was deprecated.) The default behavior of `pure_callback` and `ffi_call` under `vmap` is to now raise an exception, rather than silently producing a loop. To opt in to the previous default behavior, use `vmap_method="sequential"`.



Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index db5e0f976..810e23b10 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -80,6 +80,8 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
       `raise_to_shaped_mappings`, `reset_trace_state`, `str_eqn_compact`,
       `substitute_vars_in_output_ty`, `typecompat`, and `used_axis_names_jaxpr`. Most
       have no public replacement, though a few are available at {mod}`jax.extend.core`.
+    * The `vectorized` argument to {func}`~jax.pure_callback` and
+      {func}`~jax.ffi.ffi_call`. Use the `vmap_method` parameter instead.
 
 ## jax 0.5.3 (Mar 19, 2025)
 
diff --git a/jax/_src/callback.py b/jax/_src/callback.py
index 630539ab0..25bdb801e 100644
--- a/jax/_src/callback.py
+++ b/jax/_src/callback.py
@@ -23,7 +23,6 @@ from typing import Any
 import jax
 from jax._src import config
 from jax._src import core
-from jax._src import deprecations
 from jax._src import dispatch
 from jax._src import dtypes
 from jax._src import effects
@@ -48,10 +47,6 @@ import numpy as np
 
 logger = logging.getLogger(__name__)
 
-# TODO(dfm): Remove after 6 months.
-# Added Oct 1, 2024
-deprecations.register("jax-callback-vectorized")
-
 # `pure_callback_p` is the main primitive for staging out Python pure callbacks.
 pure_callback_p = core.Primitive("pure_callback")
 pure_callback_p.multiple_results = True
@@ -83,10 +78,9 @@ def pure_callback_impl(
     result_avals,
     callback: _FlatCallback,
     sharding: SingleDeviceSharding | None,
-    vectorized: bool | DeprecatedArg,
     vmap_method: str | None,
 ):
-  del sharding, vectorized, vmap_method, result_avals
+  del sharding, vmap_method, result_avals
   try:
     cpu_device, *_ = jax.local_devices(backend="cpu")
   except RuntimeError as e:
@@ -114,10 +108,9 @@ def pure_callback_abstract_eval(
     callback: _FlatCallback,
     result_avals,
     sharding: SingleDeviceSharding | None,
-    vectorized: bool | DeprecatedArg,
     vmap_method: str | None,
 ):
-  del avals, callback, sharding, vectorized, vmap_method
+  del avals, callback, sharding, vmap_method
   return result_avals
 
 
@@ -292,7 +285,7 @@ def pure_callback(
   When `vmap`-ed the behavior will depend on the value of the ``vmap_method``.
 
   * Calling :func:`~jax.vmap` on a callback without an explicit ``vmap_method``
-    is deprecated and it will eventually raise ``NotImplementedError``.
+    raises a ``NotImplementedError``.
   * ``vmap_method="sequential"`` uses :func:`~jax.lax.map` to loop over
     the batched arguments, calling ``callback`` once for each batch element.
   * ``vmap_method="sequential_unrolled"`` is like ``sequential``, but the loop
@@ -302,9 +295,8 @@ def pure_callback(
   * ``vmap_method="broadcast_all"`` behaves like ``expand_dims``, but the
     inputs are tiled to the expected batched shape.
 
-  If necessary, the legacy behavior provided by the deprecated
-  ``vectorized=True`` argument can be recovered using
-  ``vmap_method="legacy_vectorized"``.
+  If necessary, the legacy behavior provided by the removed ``vectorized=True``
+  argument can be recovered using ``vmap_method="legacy_vectorized"``.
 
   The current default behavior is to use ``vmap_method="sequential"`` when
   not specified, but this behavior is deprecated, and in the future, the
@@ -373,18 +365,11 @@ def pure_callback(
 
   .. _External Callbacks: https://docs.jax.dev/en/latest/external-callbacks.html
   """
-  if not isinstance(vectorized, DeprecatedArg) and not vectorized is None:
-    deprecations.warn(
-        "jax-callback-vectorized",
-        "The vectorized argument of jax.pure_callback is deprecated and setting "
-        "it will soon raise an error. To avoid an error in the future, and to "
-        "suppress this warning, please use the vmap_method argument instead.",
-        stacklevel=2)
-    if vmap_method is not None:
-      raise ValueError(
-          "the vectorized and vmap_method arguments of jax.pure_callback cannot "
-          "be used together. Please use the vmap_method argument.")
-    vmap_method = "legacy_vectorized" if vectorized else "sequential"
+  # TODO(danfm): Remove this check 3 months after v0.6.0 is released.
+  if not isinstance(vectorized, DeprecatedArg):
+    raise ValueError(
+        "The 'vectorized' argument of jax.pure_callback was removed in JAX "
+        "v0.6.0. Use 'vmap_method' instead.")
   allowed_vmap_methods = ["sequential", "sequential_unrolled", "expand_dims",
                           "broadcast_all", "legacy_vectorized", None]
   if vmap_method not in allowed_vmap_methods:
@@ -402,7 +387,6 @@ def pure_callback(
       callback=_FlatCallback(callback, in_tree),
       result_avals=tuple(flat_result_avals),
       sharding=sharding,
-      vectorized=vectorized,
       vmap_method=vmap_method,
   )
   return tree_util.tree_unflatten(out_tree, out_flat)
diff --git a/jax/_src/ffi.py b/jax/_src/ffi.py
index d1ad26754..22d54d399 100644
--- a/jax/_src/ffi.py
+++ b/jax/_src/ffi.py
@@ -24,7 +24,6 @@ import numpy as np
 
 import jax
 from jax._src import core
-from jax._src import deprecations
 from jax._src import dispatch
 from jax._src import effects
 from jax._src import util
@@ -328,7 +327,7 @@ def ffi_call(
     input_output_aliases: dict[int, int] | None = ...,
     custom_call_api_version: int = ...,
     legacy_backend_config: str | None = ...,
-    vectorized: bool | DeprecatedArg = ...,
+    vectorized: bool | None | DeprecatedArg = DeprecatedArg(),
 ) -> Callable[..., Array]:
   ...
 
@@ -345,7 +344,7 @@ def ffi_call(
     input_output_aliases: dict[int, int] | None = ...,
     custom_call_api_version: int = ...,
     legacy_backend_config: str | None = ...,
-    vectorized: bool | DeprecatedArg = ...,
+    vectorized: bool | None | DeprecatedArg = DeprecatedArg(),
 ) -> Callable[..., Sequence[Array]]:
   ...
 
@@ -361,7 +360,7 @@ def ffi_call(
     input_output_aliases: dict[int, int] | None = None,
     custom_call_api_version: int = 4,
     legacy_backend_config: str | None = None,
-    vectorized: bool | DeprecatedArg = DeprecatedArg(),
+    vectorized: bool | None | DeprecatedArg = DeprecatedArg(),
 ) -> Callable[..., Array | Sequence[Array]]:
   """Call a foreign function interface (FFI) target.
 
@@ -422,18 +421,11 @@ def ffi_call(
     to execute the FFI handler. Any keyword arguments are passed as named
     attributes to the FFI handler using XLA's FFI interface.
   """
-  if not isinstance(vectorized, DeprecatedArg) and not vectorized is None:
-    deprecations.warn(
-        "jax-callback-vectorized",
-        "The vectorized argument of ffi_call is deprecated and setting "
-        "it will soon raise an error. To avoid an error in the future, and to "
-        "suppress this warning, please use the vmap_method argument instead.",
-        stacklevel=2)
-    if vmap_method is not None:
-      raise ValueError(
-          "the vectorized and vmap_method arguments of ffi_call cannot "
-          "be used together. Please use the vmap_method argument.")
-    vmap_method = "legacy_vectorized" if vectorized else "sequential"
+  # TODO(danfm): Remove this check 3 months after v0.6.0 is released.
+  if not isinstance(vectorized, DeprecatedArg):
+    raise ValueError(
+        "The 'vectorized' argument of jax.ffi.ffi_call was removed in JAX "
+        "v0.6.0. Use 'vmap_method' instead.")
   allowed_vmap_methods = ["sequential", "sequential_unrolled", "expand_dims",
                           "broadcast_all", "legacy_vectorized", None]
   if vmap_method not in allowed_vmap_methods:
@@ -511,7 +503,6 @@ def ffi_call(
     results = ffi_call_p.bind(
         *args,
         result_avals=result_avals,
-        vectorized=vectorized,
         vmap_method=vmap_method,
         target_name=target_name,
         has_side_effect=has_side_effect,
@@ -665,21 +656,10 @@ def ffi_batching_rule(
     args,
     dims,
     *,
-    vectorized: bool | None | DeprecatedArg,
     vmap_method: str | None,
     result_avals: Sequence[core.ShapedArray],
     **kwargs: Any,
 ):
-  if isinstance(vectorized, DeprecatedArg) and vmap_method is None:
-    deprecations.warn(
-        "jax-callback-vectorized",
-        f"The default behavior of {prim.name} under vmap will soon "
-        "change. Currently, the default behavior is to generate a sequential "
-        "vmap (i.e. a loop), but in the future the default will be to raise "
-        "an error. To keep the current default, set vmap_method='sequential'.",
-        stacklevel=6)
-    vmap_method = "sequential"
-
   axis_size, = {a.shape[d] for a, d in zip(args, dims)
                 if d is not batching.not_mapped}
   new_args = [arg if dim is batching.not_mapped else
@@ -707,7 +687,6 @@ def ffi_batching_rule(
           for layout, d in zip(kwargs["input_layouts"], dims))
     outvals = prim.bind(
         *new_args,
-        vectorized=vectorized,
         vmap_method=vmap_method,
         result_avals=batched_result_avals,
         **kwargs,
@@ -723,7 +702,6 @@ def ffi_batching_rule(
           for layout in kwargs["input_layouts"])
     outvals = prim.bind(
       *bcast_args,
-      vectorized=vectorized,
       vmap_method=vmap_method,
       result_avals=batched_result_avals,
       **kwargs,
@@ -736,7 +714,6 @@ def ffi_batching_rule(
       return prim.bind(
           *merged_args,
           result_avals=result_avals,
-          vectorized=vectorized,
           vmap_method=vmap_method,
           **kwargs,
       )


Test patch:
diff --git a/tests/ffi_test.py b/tests/ffi_test.py
index abd800966..978415194 100644
--- a/tests/ffi_test.py
+++ b/tests/ffi_test.py
@@ -200,14 +200,6 @@ class FfiTest(jtu.JaxTestCase):
       else:
         self.assertArraysEqual(a, b)
 
-  @jtu.run_on_devices("gpu", "cpu")
-  def test_vectorized_deprecation(self):
-    x = self.rng().randn(3, 5, 4).astype(np.float32)
-    with self.assertWarns(DeprecationWarning):
-      ffi_call_geqrf(x, vectorized=True)
-    with self.assertWarns(DeprecationWarning):
-      jax.vmap(ffi_call_geqrf)(x)
-
   def test_input_output_aliases(self):
     def fun(x):
       return jax.ffi.ffi_call("test", x, input_output_aliases={0: 0})(x)
diff --git a/tests/python_callback_test.py b/tests/python_callback_test.py
index a8442b4a1..e4aeb8d66 100644
--- a/tests/python_callback_test.py
+++ b/tests/python_callback_test.py
@@ -1041,26 +1041,11 @@ class PureCallbackTest(jtu.JaxTestCase):
   def test_vmap_method_raise(self):
     @jax.vmap
     def f(x):
-      # Setting vectorized to None disables the current default behavior of
-      # falling back on sequential.
-      return jax.pure_callback(np.sin, x, x, vectorized=None)
+      return jax.pure_callback(np.sin, x, x)
 
     with self.assertRaisesRegex(NotImplementedError, "vmap is only supported"):
       f(jnp.arange(4.))
 
-  def test_deprecated_vectorized(self):
-    def f(x, **kwargs):
-      return jax.pure_callback(np.sin, x, x, **kwargs)
-
-    with self.assertWarnsRegex(DeprecationWarning, "The default behavior"):
-      jax.vmap(f)(jnp.arange(4.0))
-
-    with self.assertWarnsRegex(DeprecationWarning, "The vectorized argument"):
-      f(jnp.arange(4.0), vectorized=True)
-
-    with self.assertWarnsRegex(DeprecationWarning, "The vectorized argument"):
-      f(jnp.arange(4.0), vectorized=False)
-
   def test_vmap_method_expand_dims(self):
     def callback(x, y):
       self.assertTupleEqual(x.shape, (4,))