Custom_jvp and higher order differentiation

[rmoyard]

Hi Jax team, I am reworking the Jax interface for PennyLane a quantum computing framework. I am registering custom JVPs function for evaluating gradient, jacobian for techniques such as finite differences and parameter shift. For first derivative everything is working like expected. I am hitting some issues when I try to implement an higher order derivative, e.g. Hessian. I am using recursion to define the n-th derivative rule, here you can find a pseudo code of the implementation:

def execute(params, tapes, execute_fn, _n, max_diff):
    @jax.custom_jvp
    def exec(params):
        res, _ = execute_fn(tapes, params)
        return res

    @exec.defjvp
    def exec_jvp(primals, tangents):
        if _n == max_diff: 
            jvp_tapes, processing_fn = jvp(tangents)
            jvps = processing_fn(execute_fn(jvp_tapes))

        else:
            jvp_tapes, processing_fn = jvp(tangents, gradient_fn)
            
            jvps = processing_fn(
                execute(
                    jvp_tapes,
                    execute_fn,
                    _n=_n + 1,
                    max_diff=max_diff,
                )
            )

        return exec(primals), jvps

When I try to take the hessian jax.hessian(circuit, argnums = [0, 1, ...])(params, max_diff=2) : I get the following error.

jax.interpreters.ad.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.

Could you tell me what is a closed-over value? Are those tracers that were used for backprop and that cannot be differentiated over? Is there an existing example on how to solve this kind of issue.

Thank you all in advance!