Repository: google/jax

Issue:
support consts in custom batching rules
To save memory in a setup where `jax.checkpoint` was insufficient, I implemented a custom transposition via `jax.custom_derivatives.linear_call`. This worked just fine and yielded the desired savings. However, when `vmap`ing, I hit a road blocker. Unfortunately, `jax.custom_derivative.linear_call` does not implement a batching rule and thus naively applying `vmap` does not work even if the arguments to `linear_call` do support batching. Furthermore, in my case it is not possible to implement the batching manually via `jax.custom_batching.custom_vmap` because `custom_vmap` [does not work with constants in jaxpr](https://github.com/google/jax/blob/3a837c8/jax/_src/custom_batching.py#L69).

I think the best approach would be to allow both `linear_call` to be batched and `custom_vmap` to work with constants. Either of the two would unblock me :)

Below is a minimal reproducible example:
```python
from functools import partial

import jax
from jax.custom_batching import sequential_vmap
from jax.custom_derivatives import linear_call
import numpy as np


def _mul(residual_args, a, *, c):
    b, = residual_args
    c = np.array(c)  # needs to be known at trace-time
    return a * b * c


def _mul_T(residual_args, out, *, c):
    b, = residual_args
    c = np.array(c)  # needs to be known at trace-time
    return out * b * c


def mul(a, b, c):
    print(a.shape)
    return linear_call(partial(_mul, c=c), partial(_mul_T, c=c), (b, ), a)


a, b, c = np.arange(12, dtype=float), 10, np.array([2., 4.]).reshape(2, 1)

m = partial(mul, b=b, c=c)
jax.vmap(sequential_vmap(m), in_axes=(0, ))(a)
```


Discussion:
Currently, this can not be circumvented by using `jax.custom_transpose` since it is missing a batching rule too.
The current plan is to push ahead with `custom_vmap` and perhaps do away with `linear_call`. Tracking at #9073.
Ok, so in this case it would be great if `custom_vmap` would work with constants. Should I file a separate issue for that and close this one? Furthermore, to me #13298 becomes relevant then because there are some things that `linear_call` can do which `custom_transpose` can't (yet).
I'll repurpose this issue. We'll certainly need to make sure consts are supported before we consider custom batching complete. It doesn't hurt to have an open issue reminding us to do this, with your particular example registered. Thanks for filing!

Resolving patch:
diff --git a/jax/_src/custom_batching.py b/jax/_src/custom_batching.py
--- a/jax/_src/custom_batching.py
+++ b/jax/_src/custom_batching.py
@@ -19,6 +19,7 @@
 import jax
 from jax import core
 from jax import linear_util as lu
+from jax import tree_util
 from jax.interpreters import ad
 from jax.interpreters import batching
 from jax.interpreters.batching import not_mapped
@@ -66,11 +67,11 @@ def __call__(self, *args, **kwargs):
     in_avals = [core.raise_to_shaped(core.get_aval(x)) for x in args_flat]
     debug = pe.debug_info(self.fun, in_tree, False, "custom_vmap")
     jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(flat_fun, in_avals, debug)
-    assert not len(consts)
     closed_call = core.ClosedJaxpr(pe.convert_constvars_jaxpr(jaxpr), ())
+    in_tree = treedef_tuple((tree_structure(consts), in_tree))
     out_flat = custom_vmap_p.bind(*consts, *args_flat,
                                   call=closed_call,
-                                  rule=self.vmap_rule,
+                                  rule=ClosedRule(self.vmap_rule),
                                   in_tree=in_tree,
                                   out_tree=out_tree())
     return tree_unflatten(out_tree(), out_flat)
@@ -78,6 +79,21 @@ def __call__(self, *args, **kwargs):
 
 ### utils
 
+# Define a class, instead of making a function closing over `rule`, so
+# that we can override __str__
+class ClosedRule:
+  def __init__(self, rule):
+    functools.update_wrapper(self, rule)
+    self.rule = rule
+
+  def __call__(self, axis_size, all_in_batched, *all_args):
+    _, args = all_args
+    consts_batched, in_batched = all_in_batched
+    assert not any(tree_util.tree_leaves(consts_batched)), consts_batched
+    return call_rule(self.rule, axis_size, in_batched, args)
+
+  def __str__(self):
+    return str(self.rule)
 
 def ensure_list(xs):
   return xs if type(xs) is list else list(xs)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -8805,6 +8805,28 @@ def rule(axis_size, in_batched, xs):
     ys = api.vmap(f)(xs)
     self.assertAllClose(ys, jnp.cos(xs))
 
+  def test_closure(self):
+    z = jnp.array([2., 1., 3.])
+
+    @api.custom_vmap
+    def f(x): return z + jnp.sin(x)
+
+    @f.def_vmap
+    def rule(axis_size, in_batched, *args):
+      self.assertEqual(len(in_batched), 1)
+      self.assertEqual(len(args), 1)
+      xs, = args
+      xs_batched, = in_batched
+      self.assertEqual(xs_batched, True)
+      self.assertEqual(axis_size, xs.shape[0])
+      return z + jnp.cos(xs), xs_batched
+
+    x, xs = jnp.array(1.), jnp.arange(3)
+    y = f(x)
+    self.assertAllClose(y, z + jnp.sin(x))
+    ys = api.vmap(f)(xs)
+    self.assertAllClose(ys, z + jnp.cos(xs))
+
   def test_rule_multi_output(self):
     @api.custom_vmap
     def f(x): return jnp.sin(x), jnp.cos(x)
@@ -8962,6 +8984,36 @@ def rule(axis_size, in_batched, xs):
     self.assertAllClose(ys, jnp.cos(xs))
     self.assertAllClose(tys, -jnp.sin(xs) * txs)
 
+  def test_jvp_closure(self):
+    z = jnp.array([2., 1., 3.])
+    def bcast(x): return z + x - z
+
+    @api.custom_vmap
+    def f(x): return z + jnp.sin(x)
+
+    @f.def_vmap
+    def rule(axis_size, in_batched, xs):
+      self.assertEqual(axis_size, 3)
+      self.assertEqual(in_batched, [True])
+      return z + jnp.cos(xs), in_batched[0]
+
+    f_jvp = lambda x, tx: api.jvp(f, [x], [tx])
+
+    x, tx = jnp.array(1.), jnp.array(2.)
+    xs, txs = jnp.arange(3.), jnp.arange(3.) * 2.
+
+    y, ty = f_jvp(x, tx)
+    self.assertAllClose(y, z + jnp.sin(x))
+    self.assertAllClose(ty, bcast(jnp.cos(x)) * tx)
+
+    ys, tys = api.vmap(f_jvp)(xs, txs)
+    self.assertAllClose(ys, z + jnp.cos(xs))
+    self.assertAllClose(tys, bcast(-jnp.sin(xs)) * txs)
+
+    ys, tys = api.jvp(api.vmap(f), [xs], [txs])
+    self.assertAllClose(ys, z + jnp.cos(xs))
+    self.assertAllClose(tys, bcast(-jnp.sin(xs)) * txs)
+
   def test_jvp_nary(self):
     @api.custom_vmap
     def f(x, y): return jnp.sin(x) + y