Root finding options in JAX with adaptive bounds where .grad and .vmap work?

[TomClarkMassSpec]

Hi,
I've written a few versions of a code block that all work to find roots, but I keep running into trouble when using vmap or grad on them. I do see jaxopt has bisection method but that is limited. In this case I have a reasonably well behaved function where there is a single root that is positive real (of all the roots), but it does have the issue that shortly after the root the function rises quickly to numerical infinity. So I can set the lower bound at 0 but the upper bound has to be adaptive. I will share one of the many code versions that works in python on single data point but in this case fails on application of vmap due to abstract tracer value due to the boolean:

def find_positive_root_by_brentQ(function_to_solve, lowerbound, upperbound, max_iter=100, tolerance=1e-8):
    f_a = function_to_solve(lowerbound)
    fb = function_to_solve(upperbound)
    if jnp.sign(f_a) == jnp.sign(fb):
        # Expand the interval until the function values have different signs
        interval_expansion = 1.2  # Adjust the expansion factor as needed
        while jnp.sign(f_a) == jnp.sign(fb):
            upperbound *= interval_expansion
            fb = function_to_solve(upperbound)
            if upperbound > 1e20:
                return 0.0  # Unable to find a root, return default value
    if jnp.isinf(fb) or jnp.isnan(fb):
        # Adaptively reduce b until the function value is not infinity or NaN
        reduction_factor = 0.5  # Adjust the reduction factor as needed
        while jnp.isinf(fb) or jnp.isnan(fb):
            upperbound *= reduction_factor
            fb = function_to_solve(upperbound)
    for _ in range(max_iter):
        mean_target = (lowerbound + upperbound) / 2
        fc = function_to_solve(mean_target)
        if jnp.abs(fc) < tolerance:
            return mean_target
        if jnp.sign(fc) != jnp.sign(f_a):
            upperbound = mean_target
            fb = fc
        else:
            lowerbound = mean_target
            f_a = fc
        if jnp.isinf(fb) or jnp.isnan(fb):
            # Adaptively reduce b until the function value is not infinity or NaN
            while jnp.isinf(fb) or jnp.isnan(fb):
                upperbound *= reduction_factor
                fb = function_to_solve(upperbound)
    return 0.0  # Unable to find a root within the maximum number of iterations, return default value

Any suggestions for a strategy to approach this problem that's better than writing from the ground up would be appreciated.

[TomClarkMassSpec]

I'm now working on a solution that utilizes jax.scipy.optimize.minimize however the documentation states, 'minimize supports jit() compilation. It does not yet support differentiation or arguments in the form of multi-dimensional arrays, but support for both is planned.'

[mattjj]

Thanks for the question!

Can you see if jaxopt is a better fit for your needs? It's more actively developed than jax.scipy.optimize, and we should probably just delete jax.scipy.optimize and point people to that.

[TomClarkMassSpec]

Looks like there was a misunderstanding here. Jumping to the conclusion:

• Correct use of jaxopt: you have a code block built on jax. You can wrap the code block in optimizers from jaxopt.
• Incorrect use of jaxopt: you are building a code block where you want to use optimizers and expect them to work under jax transforms (grad, jit, vmap), but that doesn't work.

[patrick-kidger]

Bit late, but you might find Optimistix interesting. This is our latest scientific computing libary in JAX, and includes tackling root finding.