custom_jvp fails due to mismatched primal/output dtype with JAX 0.4.34

[YigitElma]

Hello!
I have a function that uses custom_jvp with nondiff_argnums and jit with static_argnums. The function has the derivative order as one of the inputs and the derivative with respect to the second argument is 0. A simplified version of this function is as follows,

@functools.partial(jax.custom_jvp, nondiff_argnums=(2,))
@functools.partial(jax.jit, static_argnums=(2,))
def inner(x, a, dx):
    if dx == 0:
        return (x + a).astype(jnp.float64)
    elif dx == 1:
        return jnp.ones_like(x, dtype=jnp.int64)
    else:
        return jnp.zeros_like(x, dtype=jnp.int64)

@inner.defjvp
def inner_jvp(dx, primals, tangents):
    x, a = primals
    x_dot, adot = tangents
    return inner(x, a, dx), inner(x, a, dx+1)*x_dot

This is called inside a much bigger function which I want to take the Jacobian of. Until JAX version 0.4.33 this was working fine, however with 0.4.34 release, I started to get the following error,

TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:
  primal int64[] with tangent int64[], expecting tangent ShapedArray(float0[])

Note: If you just call the jacobian of the inner, the error is not generated. So, the above code is just to show custom_jvp part.

First, I was suspecting the second argument (a) was causing the problem since its derivative is set to 0 in a hacky way. But even though I changed it to be nondiff_arg, I still get the same error.

I tried really hard to create a minimal reproducer, but I couldn't recreate it with a simple case. I am sorry for asking the question in this way, but can you think of anything that might cause the problem with the new release?

I believe patrick-kidger/equinox#871 is related to this, but this package is using an even more custom version of custom_jvp, so the implementation to my case was not straight-forward :D

Thanks!

[patrick-kidger]

The problem is your use of inner for the returned tangent. Until JAX 0.4.33 then nondifferentiable dtypes (like integers) were expected to return a tangent of the same type. From JAX 0.4.33 they ar expected to return a tangent of dtype float0. (I'm just reverse-engineering from the problem here, someone from the JAX team might be able to state this more authorativeiy.)

So in this case you have inner returning a dtype=jnp.int64, which is now no longer valid.

Given that the only meaningful tangent for a nondifferentiable type is zero, then for such any integer-dtyped-outputs I think it should suffice for you to change your returned tangent to:

return jax.custom_derivatives.SymbolicZero.from_primal_value(x)

Or if you need cross-JAX-version compatibility then you should be able to instead use:

aval = jax.core.raise_to_shaped(jax.core.get_aval(x))
if hasattr(aval, "to_tangent_aval"):
    aval = aval.to_tangent_aval()
return jax.custom_derivatives.SymbolicZero(aval)

[YigitElma]

@patrick-kidger Thank you very much for the quick response! I think my previous example was misleading, because I checked my original function and it can never return an integer. The only way I could create a similar error message with a minimal function was adding those type arguments...

The other function that is called inside the main function with custom_jvp is jax.lax.custom_root. Before commenting on that further, I will try to debug more.

I wanted to reply to explain my mistake in the previous example.

[YigitElma]

Closing in favor of #24295. Thanks again for your help @patrick-kidger !