factor convert_closure from ode to custom_derivatives

[froystig]

Adds a closure-conversion utility to custom derivatives, and in turn to our API, useful when setting up custom derivatives for higher-order functions. This is based on a subroutine in ode, with minor modifications to generalize it for wider use.

Fixes #5222

[shoyer]

This is great! One thing I might add to the docs is that you can get a pytree compatible function with tree_util.Partial(*closure_convert(objective_fn, x0)).

Another good use case for this helper would be lax.custom_root, which with closure_convert (replacing _initial_style_jaxpr) can now be implemented entirely with JAX public APIs.

[mattjj]

Nice! Idea for a test mentioned below.

[mattjj]

@shoyer I haven't thought through this, but I think that might not interact well with nondiff_argnums. Perhaps that can be sorted out but we might have to give it some thought and try out some examples.