Differentiating auxiliary variables

Hello @patrick-kidger, thank you very much for all your amazing libraries!
I did not manage to find an issue related to this but, currently we can only differentiate through the solution of e.g. a root-finding problem but not its auxiliary data.
Consider this example

import jax
import jax.numpy as jnp
import optimistix as optx


def compute_sqrt(x):
    def fn(y, args):
        def expensive_function(y):
            return y**2

        return expensive_function(y) - args, expensive_function(y)

    solver = optx.Newton(rtol=1e-5, atol=1e-5)

    y0 = jnp.array(1.0)
    sol = optx.root_find(fn, solver, y0, x, has_aux=True)
    sqrt_x = sol.value
    x_ = sol.aux
    return sqrt_x, x_


x = 2.0
print(compute_sqrt(x))
>>> (Array(1.4142135, dtype=float32), Array(2.000006, dtype=float32))
print(jax.jacobian(compute_sqrt)(x))
>>> (Array(0.35355338, dtype=float32), Array(0., dtype=float32))

where the second output of jax.jacobian is zero because auxiliary data is somehow considered as a fixed quantity.

Assuming that my residual function involves some expensive calculation, I would like to avoid the need of reevaluating x_ = expensive_function(sqrt_x) to get its gradient.

Is there a cleverer way to do this ?

Start from the root_find entry point and its has_aux handling, using the compute_sqrt example in the issue to trace how auxiliary data is treated during differentiation. The work is done when jacobian(compute_sqrt) returns the auxiliary output's dependence on x without requiring the expensive function to be reevaluated.