vmap(custom_jvp) does not strip zeros from nondifferentiable return values, leading to AD crashes

[patrick-kidger]

Description

This:

import jax
import jax.numpy as jnp

@jax.custom_jvp
def returns_int(x):
   return x, 1

@returns_int.defjvp
def f_jvp(p, t):
   (p,) = p
   (t,) = t
   p_out, p_aux = returns_int(p)
   t_out = jnp.ones_like(p_out)
   t_aux = jax.custom_derivatives.zero_from_primal(p_aux, symbolic_zeros=True)
   return (p_out, p_aux), (t_out, t_aux)

def run(x):
   _, integer = jax.vmap(returns_int)(x)
   print(integer)

jax.jvp(run, (jnp.arange(2.),), (jnp.arange(2.),))

prints:

Traced<ShapedArray(int32[2], weak_type=True)>with<JVPTrace> with
  primal = Array([1, 1], dtype=int32, weak_type=True)
  tangent = array([(b'',), (b'',)], dtype=[('float0', 'V')])

In contrast if the jax.vmap is removed, then there is no JVPTrace at all -- it seems that custom_jvp attempts to strip symbolic zeros from nondifferentiable return values, but that this is defeated by having a vmap wrapper.

So (a) there is a discrepancy there, but (b) this is now a problem for downsteam nondifferentiable primitives! They see an AD tracer, they don't have an AD rule, they explode. And moreover this spurious tangent isn't removable with lax.stop_gradient because that skips over nondifferentiable types:

jax/jax/_src/lax/lax.py

Lines 2047 to 2054 in 54fd738

	elif (dtypes.issubdtype(_dtype(x), np.floating) or
	dtypes.issubdtype(_dtype(x), np.complexfloating)):
	# break abstractions to support legacy leaked tracer use cases
	if isinstance(x, ad.JVPTracer):
	return stop(x.primal)
	return ad_util.stop_gradient_p.bind(x)
	else:
	return x

!

I imagine the fix is either to allow lax.stop_gradient to operate on nondifferentiable types, or to adjust custom_jvp to have consistent behavior regardless of whether it is vmapped. (Or maybe both?)

System info (python version, jaxlib version, accelerator, etc.)

JAX 0.4.38

[dfm]

Thanks for tracking this down and for the clear reproduction, @patrick-kidger! I haven't had a chance to dig in too deeply yet, but I think the key place where this issue hits is here:

jax/jax/_src/interpreters/batching.py

Lines 863 to 865 in 51b9fe3

	# TODO(mattjj,frostig): instantiating any SymbolicZero output is easy, but can
	# be wasteful in the rare case it actually triggers; handle symbolically!
	outs = [instantiate(replace_rule_output_symbolic_zeros(x)) for x in outs]

which has an existing TODO for @mattjj and @froystig. To get consistent behavior under vmap, I think we would need to remove the instantiate and properly handle the symbolic zeros below.

I'll take a look soon!

[rajasekharporeddy]

Hi @patrick-kidger

Testing the provided repro on Colab with JAX 0.5.0 does not result in a crash; instead, it produces the following output.

[1 1]
(None, None)

Please find the colab gist for reference.

Thank you.

[patrick-kidger]

Ah, interesting. Looks like things might have gotten fixed. I'll mention this in our upstream issue and see if things have been fixed there too.

[dfm]

Whoops - sorry I totally dropped the ball on this one. I'm glad to see that it seems to be fixed (no thanks to me!). Let me know if we should close this issue or if there are further outstanding problems.

[patrick-kidger]

Closing as per patrick-kidger/optimistix#104 (comment), looks like this is fixed! :D