Custom JVP rule for Bisection algorithm

[SNMS95]

I was trying to implement a bisection algorithm and wanted to create a custom_jvp to have more control. The code is as follows:

from jaxopt import Bisection

@partial(jax.custom_jvp, nondiff_argnums=(0,2,3))
def find_root(f, des, a, b): # a,b are the lower and upper bounds; des is an argument for f
    root = Bisection(f, a, b, check_bracket=False,
                  jit=True).run(design=des).params
    return root

@find_root.defjvp
def find_root_jvp(f, a, b, primals, tangents):
    des, = primals
    des_dot, = tangents
    root = find_root(f, des, a, b)
    # find gradient
    _, grad = jax.value_and_grad(f, argnums=(0,1))(root, des)
    rhs = -1* grad[0] / grad[1]
    root_dot = np.dot(rhs, des_dot)
    return root, root_dot

I am using the Bisection implementation from Jaxopt but it does not matter. When I test it, I do get gradients but if I check it using check_grads, it shows that the gradients are wrong.

def wrapper(vol_frac):
    def opt_func(x , design):
        return jax.nn.sigmoid(design + x).mean() - vol_frac
    return opt_func # f => optimality function i.e. f(x, design) = 0 at x = root

def dummy_objective(design, reqd_vol_frac):
    lb, ub = _construct_bounds_bisection(reqd_vol_frac, design)
    f = wrapper(reqd_vol_frac)
    root = find_root(f, design, lb, ub)
    return root

print(dummy_objective(np.array([0.5, 0.1, 0.2]), 0.5)) # Gives -0.27
print(jax.grad(dummy_objective)(np.array([0.5, 0.1, 0.2]), 0.5)) # [-3.01, -2.99, -2.98]
check_grads(lambda design: dummy_objective(design, 0.5),
                                 (np.array([0.5, 0.1, 0.2]), ), order=2, eps=1e-4) # Assertion Error!!

I tried enabling float64 but the error still persists.
Can someone tell me why this is the case and how to correct it?

[jakevdp]

I'm having trouble running your code. What is _construct_bounds_bisection? It would help when asking questions like this if you could include a minimal reproducible example.

[SNMS95]

Hi Jake, sorry for not including the _construct_bounds_bisection function when in fact the problem may be because of it.

def _construct_bounds_bisection(reqd_vol_frac: float,
                                design) -> tuple[float, float]:
    def transform_vf(vf):
        vf = np.clip(vf, 0, 1)
        return np.log(vf) - np.log1p(-vf)
    lower_bound = transform_vf(reqd_vol_frac) - np.max(design)
    upper_bound = transform_vf(reqd_vol_frac) - np.min(design)
    return lower_bound, upper_bound

I tried to condense it further but I don't know where the gradients are "leaking" because these functions are interconnected through the design variable. e.g. the optimality function (f), the lower and upper bounds (a, b) and the root all depend on design.
In pseudo-code, it may be simpler:

1. Get a 3d voxel image (Gray scale) = design (=des)
2. Find a root such that the average of this design = vol_frac (=reqd_vol_frac)
   2a. Use bisection to do this
   2b. Find the bounds using the maximum and minimum values of the design itself
3. Calculate the modified design as sigmoid(design + root)

I hope this is clearer now.

[SNMS95]

This code can be run easily and is documented. Hopefully this helps!

from jax import config
config.update("jax_enable_x64", True)
from jaxopt import Bisection
import jax
import jax.numpy as np
from functools import partial

################################################# CUSTOM JVP ##########################################
@partial(jax.custom_jvp, nondiff_argnums=(0,2,3))
def find_root(f, des, a, b):
    """Find the root of function f between a and b using bisection method."""
    root = Bisection(f, a, b, check_bracket=False,
                  jit=True).run(design=des).params
    return root

@find_root.defjvp
def find_root_jvp(f, a, b, primals, tangents):
    des, = primals
    des_dot, = tangents
    root = find_root(f, des, a, b)
    # find gradient
    _, grad = jax.value_and_grad(f, argnums=(0,1))(root, des)
    rhs = -1* grad[0] / grad[1]
    root_dot = np.dot(rhs, des_dot)
    return root, root_dot

#########################################################################################################
def wrapper(vol_frac):
    """Return the objective function for bisection method."""

    def opt_func(x , design):
        """Objective / optimality function for bisection method."""
        return jax.nn.sigmoid(design + x).mean() - vol_frac
    
    return opt_func # f(x*, design) = 0


def _construct_bounds_bisection(reqd_vol_frac,
                                design):
    """Find the bounds for bisection method."""
    def transform_vf(vf):
        vf = np.clip(vf, 0, 1)
        return np.log(vf) - np.log1p(-vf)
    
    lower_bound = transform_vf(reqd_vol_frac) - np.max(design)
    upper_bound = transform_vf(reqd_vol_frac) - np.min(design)
    return lower_bound, upper_bound


def dummy_objective(design, reqd_vol_frac):
    """A scalar valued objective function for testing."""
    lb, ub = _construct_bounds_bisection(reqd_vol_frac, design)
    f = wrapper(reqd_vol_frac)
    root = find_root(f, design, lb, ub)
    return root

######################################Testing#######################################################
design = np.array([0.5, 0.1, 0.2])
reqd_vol_frac = 0.5
print(dummy_objective(design, reqd_vol_frac))
print(jax.grad(dummy_objective)(np.array([0.5, 0.1, 0.2]), 0.5))

from jax.test_util import check_grads
check_grads(lambda design: dummy_objective(design, 0.5),
             (design, ), order=2, eps=1e-4)

[anh-tong]

There could be two points causing the wrong result:

1. This line should change from

rhs = -1* grad[0] / grad[1]

into

rhs = -1* grad[1] / grad[0]

Since grad[0] is $\partial$f/ $\partial$ root and grad[1] is $\partial$f/ $\partial$ des, it should be other way around.

See more here in jax.lax.custom_root

https://github.com/google/jax/blob/f498442daa927aeb5bdbf840bc28cc527f47d7f1/jax/_src/lax/control_flow/solves.py#L136-L143

2. jaxopt.Bisection may need to use more max_iter, and smaller tol to have more accurate root. I find this works

root = Bisection(f, a, b, check_bracket=False, maxiter=1000, tol=1e-8,
                  jit=True).run(design=des).params

[SNMS95]

Thanks a lot @jakevdp and @anh-tong.
In the end, it was a mixture of not enabling float64, tolerance for finding the root, and the tolerance set for check_grads .

[patrick-kidger]

If this is still outstanding, you might find optimistix.Bisection interesting. (GitHub link)

It also includes a custom JVP rule, for differentiating via the implicit function theorem.