Implicitly differentiate the KKT conditions

Hi,

I am currently learning how to use Jaxopt and am trying to adapt the code below to utilize its features. This code is originally from research [paper](https://arxiv.org/pdf/1906.08707.pdf) and was ported from [PyTorch ](https://github.com/locuslab/lml/blob/master/lml.py) to [Jax](https://github.com/amdee/lml/blob/lml-jax/lml.py). In the forward method, a root-finding problem is solved using the bracketing method. Meanwhile, the backward method relies on implicit differentiation through the KKT conditions.

Currently, I am only making use of Jaxopt's root-finding capabilities.
I have a specific question: How can I employ Jaxopt to perform differentiation through the KKT conditions?

Please find my current implementation below:"
```Python
import jax
import jax.numpy as jnp
import flax.linen as nn
from jaxopt import Bisection
import numpy as np

@jax.custom_vjp
def LML_jax(x, N, eps, n_iter, branch=None, verbose=0):
y, res = lml_forward(x, N, eps, n_iter, branch, verbose)
return y, res

def f(nu, x, N):
return jnp.sum(jax.nn.sigmoid(x + nu)) - N

def lml_forward(x, N, eps, n_iter, branch, verbose):
branch = branch if branch is not None else 10 if jax.devices()[0].platform == 'cpu' else 100
nx = x.shape[0]
if nx <= N:
return jnp.ones(nx, dtype=x.dtype), None

x_sorted = jnp.sort(x)[::-1]
nu_lower = -x_sorted[N-1] - 7.
nu_upper = -x_sorted[N] + 7.

# Using Bisection from jaxopt
bisection = Bisection(optimality_fun=f, lower=nu_lower, upper=nu_upper, tol=eps, check_bracket=False)
sol = bisection.run(x=x, N=N)
nu = sol.params

y = jax.nn.sigmoid(x + nu)
return y, (y, nu, x, N)

def lml_backward(res, grad_output):
y, nu, x, N = res
if y is None:
return (jnp.zeros_like(x), None, None, None, None, None)

Hinv = 1. / (1. / y + 1. / (1. - y))
dnu = jnp.sum(Hinv * grad_output) / jnp.sum(Hinv)
dx = -Hinv * (-grad_output + dnu)
return (dx, None, None, None, None, None)

LML_jax.defvjp(lml_forward, lml_backward)

class LML(nn.Module):
N: int = 1
eps: float = 1e-4
n_iter: int = 100
branch: int = None
verbose: int = 0

@nn.compact
def __call__(self, x):
return LML_jax(x, N=self.N, eps=self.eps, n_iter=self.n_iter, branch=self.branch, verbose=self.verbose)

if __name__ == '__main__':
m = 10
n = 2
np.random.seed(0)
x = np.random.random(m)
x_jax_unbatched = jnp.array(x)
x_jax_batched = jnp.stack([x_jax_unbatched, x_jax_unbatched])
x = jnp.stack([x, x])
model = LML(N=n)
key1, key2 = jax.random.split(jax.random.PRNGKey(1))
dummy_input = jax.random.normal(key1, (n, m))
params = model.init(jax.random.PRNGKey(0), dummy_input)
LML_state = model.bind(params)
lml = lambda x_input: LML_state(x_input)[0]

y_unbatched = lml(x_jax_unbatched)
y_batched = jax.vmap(lml)(x_jax_batched)
y_unbatched_check = np.array(y_unbatched, copy=False)
y_batched_check = np.array(y_batched, copy=False)

vyo_unbatched, dyo_unbatched = jax.value_and_grad(lml)(x_jax_unbatched)
vyo_batched, dyo_batched = jax.vmap(jax.value_and_grad(lml))(x_jax_batched)
print(f"value of y vyo_unbatched: {vyo_unbatched}\ngradient of y dy0 vyo_unbatched: {dyo_unbatched}")
print(f"\nvalue of y vyo_batched: {vyo_batched}\ngradient of y dy0 vyo_batched: {dyo_batched}")
```