Open jakevdp opened 4 weeks ago
I think this is actually a bug (although @mattjj and @froystig will know better, of course!). If we update f_jvp
to:
@f.defjvp
def f_jvp(indices, primals, tangents):
x, = primals
x_dot, = tangents
return x[indicies], x_dot[indices]
everything works as it should without any leaked tracers. I haven't had a chance to dig into why this is, but there's a lot of subtlety around handling these recursive patterns. Regardless, I don't think that we do need to require that "nondiff" args be "static".
OK maybe I take that back! It's clear from this error message:
that custom VJP at least has this requirement that nondiff args cannot be tracers. I think it also follows that they can't be tracers because they're handled by baking them into the functions, rather than binding them. I guess from my example above, there are some cases where things might still work, but they shouldn't be expected to.
Perhaps it would be reasonable to add the tracer check from custom_vjp
to custom_jvp
?
Edited to add: These are two relevant tests:
and it looks like we currently expect the custom_jvp
to work with some tracers (BatchTracer
) but not others. Reading the PR adding that comment https://github.com/google/jax/pull/14263, I think perhaps it makes sense to add that check to custom_jvp
, but perhaps there are people depending on the current more relaxed behavior.
The problem is here that
indices
is a traced array, but is marked as anondiff_argnum
. This is user error, but we should fail with a more informative error.