Twice differentiability of OSQP

Hi
Thank you for this very cool repo and, sorry, I'm new to "Issues" and to jax,

I'm getting an error of `JaxStackTraceBeforeTransformation` when trying to take the derivative with a `ProjectedGradient` solver involved. The solver itself runs flawlessly and derivatives seem to work without the solver.
I want to later use `outer_with_proj` (with a slightly different objective) for implicit differentiation of `theta`.

Any advice?

**The code:**
```
import numpy as np
import jax.numpy as jnp
import jax
import random
import matplotlib.pyplot as plt

random.seed(0)

# parameters
m = 3
n = 25
nm = n*m
d = 2
d0 = 5
d1 = 5
sshape = (d0,d1,m)

# polyhedron constraints
# Az = b
A = np.zeros((n, nm))
for i in range(n):
A[i, i*m:(i+1)*m] = 1

b = np.ones(n)

# Jz <= h
J = np.zeros((2*nm, nm))
J[:nm, :] = -np.eye(nm)
J[nm:, :] = np.eye(nm)
h = np.zeros(2*nm)
h[:nm] = 0
h[nm:] = 1

# initial state
g0 = np.random.rand(n, m)
g0 = (g0 / np.sum(g0, 1, keepdims=True)).flatten()

# auxilary params for objective
L = np.stack([np.repeat(range(d0), d1), np.tile(range(d1), d0)]).T

x = L[...,0]
y = L[...,1]

library = np.stack([x**0, x, y,], axis=-1)
library = library / library.max(axis=0, keepdims=True)
l = library.shape[1]

# defining a gradient
params = np.zeros((l, m))
params[1,0] = -1
params[0,0] = 1
params[2,1] = -1
params[0,1] = 1
params[0,2] = 1

grads = (library @ params)

plt.imshow(grads.reshape(sshape))

# convert to jax
A = jnp.asarray(A)
b = jnp.asarray(b)
J = jnp.asarray(J)
h = jnp.asarray(h)
g0 = jnp.asarray(g0)

library = jnp.asarray(library)
params = jnp.asarray(params)
grads = jnp.asarray(grads.flatten())

# problem setup
from jaxopt import projection
from jaxopt import ProjectedGradient

lr = 0.001

def obj(g, theta):
"""
Maximize the sum of the log of the weighted sum of the gradients
:param g: flattened G nxm
:param theta: flattened params lxm
"""

G = jnp.reshape(g, (n,m))
w = jnp.einsum('nl,lm->nm', jnp.reshape(grads, (n,l)), jnp.reshape(theta, (l,m)))
return -jnp.sum(jnp.log(jnp.sum(jnp.multiply(w, G), 0)))

def outer_with_proj(g0, theta):
"""
Given theta, compute error from gradients where each sample is within the simplex
:param g0: flattened G nxm
:param theta: flattened params lxm
"""
def proj(p,C):
return projection.projection_polyhedron(p, C, check_feasible=False)

solver = ProjectedGradient(fun=obj,
projection=proj,
maxiter=10,
implicit_diff=True,
)

g_fit = solver.run(g0, (A, b, J, h), theta).params
return jnp.mean((g_fit - grads) ** 2)

def outer_no_proj(g0, theta):
"""
Given theta, compute error from gradients
:param g0: flattened G nxm
:param theta: flattened params lxm
"""
g_fit = g0 - lr * jax.grad(obj, argnums=0)(g0, theta)
return jnp.mean((g_fit - grads) ** 2)

eps = 1e-2
theta = params + eps
print(outer_no_proj(g0, theta))
print(outer_with_proj(g0, theta))

print(jax.grad(outer_no_proj, argnums=1)(g0, theta))
print(jax.grad(outer_with_proj, argnums=1)(g0, theta))
```

**The error:**
```
---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation Traceback (most recent call last)
File [~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/runpy.py:196](https://file+.vscode-resource.vscode-cdn.net/Users/nomo/PycharmProjects/pareto_dynamic/notebooks/~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/runpy.py:196), in _run_module_as_main(***failed resolving arguments***)
195 sys.argv[0] = mod_spec.origin
--> 196 return _run_code(code, main_globals, None,
197 "__main__", mod_spec)

File [~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/runpy.py:86](https://file+.vscode-resource.vscode-cdn.net/Users/nomo/PycharmProjects/pareto_dynamic/notebooks/~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/runpy.py:86), in _run_code(***failed resolving arguments***)
79 run_globals.update(__name__ = mod_name,
80 __file__ = fname,
81 __cached__ = cached,
(...)
84 __package__ = pkg_name,
85 __spec__ = mod_spec)
---> 86 exec(code, run_globals)
87 return run_globals

File [~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/site-packages/ipykernel_launcher.py:17](https://file+.vscode-resource.vscode-cdn.net/Users/nomo/PycharmProjects/pareto_dynamic/notebooks/~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/site-packages/ipykernel_launcher.py:17)
15 from ipykernel import kernelapp as app
---> 17 app.launch_new_instance()

File [~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/site-packages/traitlets/config/application.py:1043](https://file+.vscode-resource.vscode-cdn.net/Users/nomo/PycharmProjects/pareto_dynamic/notebooks/~/opt/miniconda3/envs/pareto_dyn/lib/python3.10/site-packages/traitlets/config/application.py:1043), in Application.launch_instance(***failed resolving arguments***)
1042 app.initialize(argv)
-> 1043 app.start()
...
372 *(_flatten(params.transpose()) + x_cotangent),
373 const_lengths=const_lengths.transpose(), jaxprs=jaxprs.transpose())
374 # drop aux values in cotangent computation

AssertionError:
```

Thanks again!