Need help for error when trying to make jet respect custom jvp rules

[y1xiaoc]

Hi,

I'm trying to modify the jet tracer so that it will respect the custom jvp rule (up to a certain order). Right now, I just need up to second order so just brute force jvp twice should work for my use case. However, after I tried to implement that in the jet.JetTrace.process_custom_jvp_call I get an error from post_process_custom_jvp_call of JVPTrace when I'm trying to do 2nd order jet for functions that have custom_jvp defined (like slogdet). (Order 1 jet works fine, which is expected as I'm following jvp trace basically)

My implementation is attached below. I borrowed some code from process_primitive and JVPTrace. I don't think I fully understood the full_lower or the linear utils wrapped function. I also failed to figured out clearly the details of the mechanism of custom_jvp. So I guess I must did something wrong for the tracer so I'm wondering if there could be a fix for the error. Thanks a lot!

def process_custom_jvp_call(self, primitive, fun, f_jvp, tracers, *,
                            symbolic_zeros):
    # copy this part from process_primitive
    order = self.main.order
    primals_in, series_in = unzip2((t.primal, t.terms) for t in tracers)
    primals_in = map(core.full_lower, primals_in) # not needed?
    series_in = [[zero_term] * order if s is zero_series else s
                 for s in series_in]
    series_in = [[jnp.zeros(np.shape(x), dtype=jnp.result_type(x))
                  if t is zero_term else t for t in series]
                 for x, series in zip(primals_in, series_in)]
    
    def my_jvp(primals_in, tangets_in):
        outs = f_jvp.call_wrapped(*primals_in, *tangets_in)
        return split_list(outs, [len(outs) // 2])
    
    z0 = tuple(primals_in)
    z1 = tuple(s[0] for s in series_in)

    if order == 1:
        o0, o1 = my_jvp(z0, z1)
        primal_out, terms_out = o0, list(zip(o1))
        return [JetTracer(self, p, ts) for p, ts in zip(primal_out, terms_out)]
    
    z2 = tuple(s[1] for s in series_in)
    
    inner_jvp = lambda *a: my_jvp(a, z1)
    (o0, o1), (_, o2_1) = jax.jvp(inner_jvp, z0, z1) # <- the error occurs here
    _, o2_2 = my_jvp(z0, z2)

    cast_f0s = lambda term: jax.tree_map(replace_float0s, o0, term)
    o1, o2_1, o2_2 = map(cast_f0s, (o1, o2_1, o2_2))
    o2 = jax.tree_map(lax.add, o2_1, o2_2)

    primal_out, terms_out = o0, list(zip(o1, o2))
    return [JetTracer(self, p, ts) for p, ts in zip(primal_out, terms_out)]

And the error is the following (only the last few frames):

.../jet_example.ipynb Cell 42 line 3
     29 # print(z0, z1, z2, sep="\n")
     31 inner_jvp = lambda *a: my_jvp(a, z1)
---> 32 (o0, o1), (_, o2_1) = jax.jvp(inner_jvp, z0, z1)
     33 # _ = jax.jvp(f_jvp.call_wrapped, (*z0, *z1), (*z1, *jax.tree_map(jnp.zeros_like, z1)))
     34 _, o2_2 = my_jvp(z0, z2)

    [... skipping hidden 3 frame]

.../jet_example.ipynb Cell 42 line 3
     28 z2 = tuple(s[1] for s in series_in)
     29 # print(z0, z1, z2, sep="\n")
---> 31 inner_jvp = lambda *a: my_jvp(a, z1)
     32 (o0, o1), (_, o2_1) = jax.jvp(inner_jvp, z0, z1)
     33 # _ = jax.jvp(f_jvp.call_wrapped, (*z0, *z1), (*z1, *jax.tree_map(jnp.zeros_like, z1)))

.../jet_example.ipynb Cell 42 line 1
     15 def my_jvp(primals_in, tangets_in):
---> 16     outs = f_jvp.call_wrapped(*primals_in, *tangets_in)
     17     return split_list(outs, [len(outs) // 2])

    [... skipping hidden 3 frame]

File.../python3.11/site-packages/jax/_src/interpreters/ad.py:389, in JVPTrace.post_process_custom_jvp_call(self, out_tracers, _)
    388 def post_process_custom_jvp_call(self, out_tracers, _):
--> 389   raise CustomJVPException()

CustomJVPException: Detected differentiation of a custom_jvp function with respect to a closed-over value. That isn't supported because the custom JVP rule only specifies how to differentiate the custom_jvp function with respect to explicit input parameters. Try passing the closed-over value into the custom_jvp function as an argument, and adapting the custom_jvp rule.