jaxopt.Bisection() does not support jax.hessian()

[yyang97]

def F(x, factor):
    return factor * x ** 3 - x - 2

def root(factor):
    bisec = jaxopt.Bisection(optimality_fun=F, lower=1, upper=2)
    return bisec.run(factor=factor).params

# hessian of root with respect to factor at 2.0.
jax.hessian(root)(2.0)
# TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.

When my project is involved with calculating the Hessian of the root finding, I just found that jaxopt.Bisection() does not support that...

Do I implement it correctly or it just does not support Hessian?

[TomClarkMassSpec]

Hi,

Not an answer but I have a similar issue.

I found (and was hoping) jaxopt.bisection solved my need for a way to find a slightly complicated root that is compatible with jax.vmap and jax.grad and next jax.hessian.

I am told:, " Most solvers in JAXopt don't work out of the box with scalars (maybe they should?) So typically, ones needs to use an array of size 1 instead. " and I need to find a work around or a bit of clarity.

Indeed jaxopt.bisection produced the correct root individually for multiple tests but later I tried to vmap a a broader code and it turns out my problem lies with data types and jaxopt.bisection.

At this stage of my code mostly everything is a jax device, dtype = jnp.float64 array including single valued jax arrays.

What was thought to work did amazingly well on a single datapoint basis:

import jax.numpy as jnp
from jax import grad, vmap, value_and_grad, lax, jit
from jaxopt import Bisection

def broader_function(custom_tree, ..., regime, ...):
    ...
    alt_pytree = (Beta1.item(), Beta2.item(), Beta3.item(), Beta4.item(), Beta5.item(), terminal_sp.item(), length.item(),regime.item())
    ...    
    lower_bound = jnp.array(0.0)
    upper_bound = jnp.array(100000) 
    t_initial = jnp.array([1.0],dtype=jnp.float64)
    tolerance = 1e-50
    # jaxopt.bisection
    optimizer_bisec_Thv = Bisection(optimality_fun = root_from_speed_by_jaxopt_bisection , lower=lower_bound, upper=upper_bound, maxiter=500, tol=tolerance, check_bracket=True, verbose=False, implicit_diff_solve=None, has_aux=False, jit='auto', unroll='auto')
    result_optimizer_bisec_Thv = optimizer_bisec_Thv.run(t_initial, alt_pytree)
    Thv = result_optimizer_bisec_Thv.params
    .....
    return Thv,....

def root_from_speed_by_jaxopt_bisection(t, alt_pytree):
    beta1, beta2, beta3, beta4, beta5, vterm, length, Regime = alt_pytree
    t_dependent_terms = (-beta2 + beta5)*jnp.exp(((-beta2 + beta5)*t)/(2*beta1)) - (beta2 + beta5)*jnp.exp(-((beta2 + beta5)*t)/(2*beta1))
    t_independent_terms = -4*vterm*beta3*beta1/beta4
    return t_dependent_terms + t_independent_terms

Turns out .item() is not compatible with jax tracing and applying .vmap or .grad. Using .item() succeeds in a single case basis. Alternatively, not using .item() and sending as jnp.array to Bisection causes boolean errors.

For single use cases, all of the following are jax device arrays with a single float 64 entry.
Beta1.item(), Beta2.item(), Beta3.item(), Beta4.item(), Beta5.item(), terminal_sp.item(), length.item(),regime.item()

When applying vmap, I am expanding over a jax array of regime values to return an array of answers.

Struggling to find a way past this apparent data type problem in using jaxopt.bisection.

Thanks,
Tom

[TomClarkMassSpec]

Half my problem is fixed. For the single case this works with removing the d-type on t_initial which now permits alt_pytree to accept 0-dim jax_arrays (scalars).

    
    alt_pytree = (Beta1, Beta2, Beta3, Beta4, Beta5, terminal_sp, length, regime)   
 
    lower_bound = jnp.array(0.0)
    upper_bound = jnp.array(100000.) 
    t_initial = jnp.array([1.0])
    tolerance = 1e-50

[TomClarkMassSpec]

assigning vmap returns error.

def calculate_hang_time_wrapper(regime, constant_args):
    # Extract the constant arguments
    (custom_tree,...) = constant_args
    
    # Call the calculate_hang_time function with the arguments and the regime
    return calculate_hang_time(custom_tree, ..., regime, ...)

hang_time_array = jax.vmap(calculate_hang_time_wrapper, in_axes=(0, None))(regime_array, constant_args)

state = self.init_state(init_params, *args, **kwargs)

File "/opt/anaconda3/lib/python3.8/site-packages/jaxopt/_src/bisection.py", line 105, in init_state
if sign == 0:

File "/opt/anaconda3/lib/python3.8/site-packages/jax/core.py", line 634, in bool
def bool(self): return self.aval._bool(self)

File "/opt/anaconda3/lib/python3.8/site-packages/jax/core.py", line 1267, in error
raise ConcretizationTypeError(arg, fname_context)

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[], weak_type=True)>with<BatchTrace(level=1/1)> with
val = DeviceArray([False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False, False, False, False, False, False,
False, False, False], dtype=bool, weak_type=True)
batch_dim = 0
The problem arose with the bool function.

[TomClarkMassSpec]

Restructured and rewrote the entire code block to apply jax.vmap to a code block containing jaxopt.bisection isn a more direct way. I still get the same concretization boolean error from within jaxopt.bisection. The Bisection solver works for a single input set and appears to not work with vmap.

[mblondel]

Have you tried to set check_bracket=False?

[mblondel]

(it would have been better to create another thread)

[mblondel]

@yyang97 Indeed, second-order derivatives via implicit diff are not supported yet in JAXopt :( I'm curious, what is your use case?

[mblondel]

In the mean time, if the parameter w.r.t. you want to differentiate is a scalar, you could use a finite-difference of the gradient, although it's not ideal.

[yyang97]

Hi Maintainer,

Thank you very much for your reply. I am using Maximum Likelihood Estimate (MLE) to estimate parameters, where the likelihood is involved with root finding jaxopt.Bisection(). To calculate the standard error of the estimator, it requires second order derivatives and hence we need second order derivatives involved with jaxopt.Bisection().