Repository: google/jax
Pull request #22869 (2024-08-05 16:56:35)

Small tweaks to custom_vmap UI
I'm working on some extensions to `custom_vmap` and came across these small UI improvements (I think!). This includes the two changes:

1. A weakening of the `kwargs` check to be consistent with the one in `custom_vjp`/`custom_jvp`, and
2. An improved error message when `def_vmap` isn't called.


Patch:
diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml
index 5d46f8fbf..743aae9b4 100644
--- a/.github/workflows/ci-build.yaml
+++ b/.github/workflows/ci-build.yaml
@@ -141,7 +141,7 @@ jobs:
         PY_COLORS: 1
       run: |
         pytest -n auto --tb=short --doctest-glob='*.md' --doctest-glob='*.rst' docs --doctest-continue-on-failure --ignore=docs/multi_process.md --ignore=docs/jax.experimental.array_api.rst
-        pytest -n auto --tb=short --doctest-modules jax --ignore=jax/config.py --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/_src/lib/triton.py --ignore=jax/_src/lib/mosaic_gpu.py --ignore=jax/interpreters/mlir.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py --ignore=jax/_src/tpu_custom_call.py --ignore=jax/experimental/mosaic --ignore=jax/experimental/pallas --ignore=jax/_src/pallas --ignore=jax/experimental/array_api --ignore=jax/lib/xla_extension.py
+        pytest -n auto --tb=short --doctest-modules jax --ignore=jax/config.py --ignore=jax/experimental/jax2tf --ignore=jax/_src/lib/mlir --ignore=jax/_src/lib/triton.py --ignore=jax/_src/lib/mosaic_gpu.py --ignore=jax/interpreters/mlir.py --ignore=jax/experimental/array_serialization --ignore=jax/collect_profile.py --ignore=jax/_src/tpu_custom_call.py --ignore=jax/experimental/mosaic --ignore=jax/experimental/pallas --ignore=jax/_src/pallas --ignore=jax/experimental/array_api
 
 
   documentation_render:
diff --git a/jax/_src/api_util.py b/jax/_src/api_util.py
index dd1cdcbe6..481dec006 100644
--- a/jax/_src/api_util.py
+++ b/jax/_src/api_util.py
@@ -556,6 +556,17 @@ def _assert_no_intersection(static_argnames, donate_argnames):
         f"{out} appear in both static_argnames and donate_argnames")
 
 
+def resolve_kwargs(fun: Callable, args, kwargs):
+  if isinstance(fun, partial):
+    fun = lambda *args, **kwargs: None
+  ba = inspect.signature(fun).bind(*args, **kwargs)
+  ba.apply_defaults()
+  if ba.kwargs:
+    raise TypeError("keyword arguments could not be resolved to positions")
+  else:
+    return ba.args
+
+
 def _dtype(x):
   try:
     return dtypes.result_type(x)
diff --git a/jax/_src/custom_batching.py b/jax/_src/custom_batching.py
index 4d41849b7..4b859e910 100644
--- a/jax/_src/custom_batching.py
+++ b/jax/_src/custom_batching.py
@@ -27,7 +27,7 @@ from jax._src import source_info_util
 from jax._src import traceback_util
 from jax._src import tree_util
 from jax._src import util
-from jax._src.api_util import flatten_fun_nokwargs
+from jax._src.api_util import flatten_fun_nokwargs, resolve_kwargs
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters.batching import not_mapped
@@ -64,7 +64,12 @@ class custom_vmap:
 
   @traceback_util.api_boundary
   def __call__(self, *args, **kwargs):
-    assert not kwargs
+    fun_name = getattr(self.fun, "__name__", str(self.fun))
+    if not self.vmap_rule:
+      raise AttributeError(
+          f"No batching rule defined for custom_vmap function {fun_name} "
+          "using def_vmap.")
+    args = resolve_kwargs(self.fun, args, kwargs)
     args_flat, in_tree = tree_flatten(args)
     flat_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(self.fun), in_tree)
     in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat]
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index d27b0efc7..bc9f7a687 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -17,7 +17,6 @@ from __future__ import annotations
 from collections.abc import Callable, Sequence
 import dataclasses
 from functools import update_wrapper, reduce, partial, wraps
-import inspect
 from typing import Any, Generic, TypeVar
 
 from jax._src import config
@@ -30,7 +29,8 @@ from jax._src import linear_util as lu
 from jax._src import traceback_util
 from jax._src.ad_util import (
     stop_gradient_p, SymbolicZero, Zero, zeros_like_aval)
-from jax._src.api_util import argnums_partial, flatten_fun_nokwargs
+from jax._src.api_util import (
+    argnums_partial, flatten_fun_nokwargs, resolve_kwargs)
 from jax._src.core import raise_to_shaped
 from jax._src.errors import UnexpectedTracerError
 from jax._src.interpreters import ad
@@ -56,17 +56,6 @@ zip = safe_zip
 
 ### util
 
-def _resolve_kwargs(fun, args, kwargs):
-  if isinstance(fun, partial):
-    # functools.partial should have an opaque signature.
-    fun = lambda *args, **kwargs: None
-  ba = inspect.signature(fun).bind(*args, **kwargs)
-  ba.apply_defaults()
-  if ba.kwargs:
-    raise TypeError("keyword arguments could not be resolved to positions")
-  else:
-    return ba.args
-
 def _initial_style_jaxpr(fun, in_avals):
   jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(fun, in_avals)
   return jaxpr, consts
@@ -240,7 +229,7 @@ class custom_jvp(Generic[ReturnValue]):
       msg = f"No JVP defined for custom_jvp function {primal_name} using defjvp."
       raise AttributeError(msg)
     jvp_name    = getattr(self.jvp, '__name__', str(self.jvp))
-    args = _resolve_kwargs(self.fun, args, kwargs)
+    args = resolve_kwargs(self.fun, args, kwargs)
     if self.nondiff_argnums:
       nondiff_argnums = set(self.nondiff_argnums)
       args = tuple(_stop_gradient(x) if i in nondiff_argnums else x
@@ -599,7 +588,7 @@ class custom_vjp(Generic[ReturnValue]):
       msg = f"No VJP defined for custom_vjp function {primal_name} using defvjp."
       raise AttributeError(msg)
     fwd_name = getattr(self.fwd, '__name__', str(self.fwd))
-    args = _resolve_kwargs(self.fun, args, kwargs)
+    args = resolve_kwargs(self.fun, args, kwargs)
     if self.optimize_remat:
       fwd = optimize_remat_of_custom_vjp_fwd(
           self.fun, self.fwd, nondiff_argnums=self.nondiff_argnums,
@@ -1451,7 +1440,7 @@ def optimize_remat_of_custom_vjp_fwd(
     # above and it would be good to consolidate it.
     primal_name = getattr(fun, "__name__", str(fun))
     fwd_name = getattr(fwd, "__name__", str(fwd))
-    args = _resolve_kwargs(fwd, args, kwargs)
+    args = resolve_kwargs(fwd, args, kwargs)
     if nondiff_argnums:
       for i in nondiff_argnums: _check_for_tracers(args[i])
       nondiff_argnums_ = set(nondiff_argnums)
diff --git a/jax/lib/__init__.py b/jax/lib/__init__.py
index 7534cd6c7..7a380348f 100644
--- a/jax/lib/__init__.py
+++ b/jax/lib/__init__.py
@@ -15,9 +15,7 @@
 # ruff: noqa: F401
 from jax._src.lib import (
   version_str as __version__,
-)
-from jax.lib import (
-  xla_bridge as xla_bridge,
   xla_client as xla_client,
   xla_extension as xla_extension,
 )
+from jax.lib import xla_bridge as xla_bridge
diff --git a/jax/lib/xla_client.py b/jax/lib/xla_client.py
deleted file mode 100644
index e88a4247a..000000000
--- a/jax/lib/xla_client.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# Copyright 2024 The JAX Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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.
-
-from jax._src.lib import xla_client as _xc
-
-_xla = _xc._xla  # TODO(jakevdp): deprecate this in favor of jax.lib.xla_extension
-bfloat16 = _xc.bfloat16  # TODO(jakevdp): deprecate this in favor of ml_dtypes.bfloat16
-
-dtype_to_etype = _xc.dtype_to_etype
-execute_with_python_values = _xc.execute_with_python_values
-get_topology_for_devices = _xc.get_topology_for_devices
-heap_profile = _xc.heap_profile
-mlir_api_version = _xc.mlir_api_version
-ops = _xc.ops
-register_custom_call_target = _xc.register_custom_call_target
-shape_from_pyval = _xc.shape_from_pyval
-ArrayImpl = _xc.ArrayImpl
-Client = _xc.Client
-CompileOptions = _xc.CompileOptions
-Device = _xc.Device
-DeviceAssignment = _xc.DeviceAssignment
-FftType = _xc.FftType
-Frame = _xc.Frame
-HloSharding = _xc.HloSharding
-OpSharding = _xc.OpSharding
-PaddingType = _xc.PaddingType
-PrimitiveType = _xc.PrimitiveType
-Shape = _xc.Shape
-Traceback = _xc.Traceback
-XlaBuilder = _xc.XlaBuilder
-XlaComputation = _xc.XlaComputation
-XlaRuntimeError = _xc.XlaRuntimeError
-
-del _xc
diff --git a/jax/lib/xla_extension.py b/jax/lib/xla_extension.py
deleted file mode 100644
index 14f57b61c..000000000
--- a/jax/lib/xla_extension.py
+++ /dev/null
@@ -1,37 +0,0 @@
-# Copyright 2024 The JAX Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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.
-
-from jax._src.lib import xla_extension as _xe
-
-get_distributed_runtime_client = _xe.get_distributed_runtime_client
-get_distributed_runtime_service = _xe.get_distributed_runtime_service
-hlo_module_cost_analysis = _xe.hlo_module_cost_analysis
-hlo_module_to_dot_graph = _xe.hlo_module_to_dot_graph
-ifrt_proxy = _xe.ifrt_proxy
-jax_jit = _xe.jax_jit
-mlir = _xe.mlir
-pmap_lib = _xe.pmap_lib
-profiler = _xe.profiler
-pytree = _xe.pytree
-ArrayImpl = _xe.ArrayImpl
-Device = _xe.Device
-DistributedRuntimeClient = _xe.DistributedRuntimeClient
-HloPrintOptions = _xe.HloPrintOptions
-OpSharding = _xe.OpSharding
-PjitFunctionCache = _xe.PjitFunctionCache
-PjitFunction = _xe.PjitFunction
-PmapFunction = _xe.PmapFunction
-XlaRuntimeError = _xe.XlaRuntimeError
-
-del _xe


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index cb0d7c0d4..4aafc42b7 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -10798,6 +10798,32 @@ class CustomVmapTest(jtu.JaxTestCase):
 
     self.assertAllClose(y, (x + a)**2)
 
+  def test_kwargs(self):
+    @jax.custom_batching.custom_vmap
+    def f(x): return jnp.sin(x)
+
+    @f.def_vmap
+    def rule(axis_size, in_batched, xs):
+      xs_batched, = in_batched
+      self.assertEqual(xs_batched, True)
+      self.assertEqual(axis_size, xs.shape[0])
+      return jnp.cos(xs), xs_batched
+
+    x, xs = jnp.array(1.), jnp.arange(3)
+    y = f(x=x)
+    self.assertAllClose(y, jnp.sin(x))
+    ys = api.vmap(f)(x=xs)
+    self.assertAllClose(ys, jnp.cos(xs))
+
+  def test_undefined_rule(self):
+    @jax.custom_batching.custom_vmap
+    def f(x): return jnp.sin(x)
+
+    with self.assertRaisesRegex(
+        AttributeError, "No batching rule defined for custom_vmap function f"):
+      f(0.5)
+
+
 class CustomApiTest(jtu.JaxTestCase):
   """Test interactions among the custom_{vmap,jvp,vjp,transpose,*} APIs"""