Help in running examples

Hi!

Thank you for the package and paper!

I am interested in applying JaxOpt for implicit differentiation applications.

I have been trying to replicate simple examples from the paper, but I seem to always get the same issue.

For instance, if I implement a ridge regression solver as in page 3 of the paper:
```
# build dataset
w = jnp.array([1, 2, 3, 4])
sigma = 2

npts = 1000

X_tr = jnp.array(np.random.randn(npts, w.shape[0]))

y_tr = jnp.dot(X_tr, w) + jnp.array(np.random.randn(npts)) * sigma

# ridge regression, paper page 3.
def f(x, theta):
residual = jnp.dot(X_tr, x) - y_tr
return (jnp.sum(residual ** 2) + theta * jnp.sum(x**2)) / 2

F = jax.grad(f)

@implicit_diff.custom_root(F)
def ridge_solver(theta):
XX = jnp.dot(X_tr.T, X_tr)
Xy = jnp.dot(X_tr.T, y_tr)
I = jnp.eye(X_tr.shape[1])
x_star = jnp.linalg.solve(XX + theta * I, Xy)
return x_star
```

I am able to compute the jacobian `F`:
```
jax.jacobian(F)(ridge_solver(10.), 10.)

Array([[ 1.03519568e+03, -1.38804913e+00, 1.67368546e+01,
-3.62912827e+01],
[-1.38804913e+00, 1.03733008e+03, -6.48666048e+00,
2.09004173e+01],
[ 1.67368546e+01, -6.48666048e+00, 1.04790186e+03,
-9.89336014e-01],
[-3.62912827e+01, 2.09004173e+01, -9.89336014e-01,
1.00473975e+03]], dtype=float32)
```
but I don't seem to be able to compute the jacobian by calling the inner solver directly:
```
jax.jacobian(ridge_solver)(10.)
---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation Traceback (most recent call last)
File ~/miniconda3/envs/default/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 ~/miniconda3/envs/default/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 ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel_launcher.py:17, in
15 from ipykernel import kernelapp as app
---> 17 app.launch_new_instance()

File ~/miniconda3/envs/default/lib/python3.10/site-packages/traitlets/config/application.py:976, in Application.launch_instance(***failed resolving arguments***)
975 app.initialize(argv)
--> 976 app.start()

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/kernelapp.py:712, in IPKernelApp.start(***failed resolving arguments***)
711 try:
--> 712 self.io_loop.start()
713 except KeyboardInterrupt:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/tornado/platform/asyncio.py:199, in BaseAsyncIOLoop.start(***failed resolving arguments***)
198 asyncio.set_event_loop(self.asyncio_loop)
--> 199 self.asyncio_loop.run_forever()
200 finally:

File ~/miniconda3/envs/default/lib/python3.10/asyncio/base_events.py:600, in BaseEventLoop.run_forever(***failed resolving arguments***)
599 while True:
--> 600 self._run_once()
601 if self._stopping:

File ~/miniconda3/envs/default/lib/python3.10/asyncio/base_events.py:1896, in BaseEventLoop._run_once(***failed resolving arguments***)
1895 else:
-> 1896 handle._run()
1897 handle = None

File ~/miniconda3/envs/default/lib/python3.10/asyncio/events.py:80, in Handle._run(***failed resolving arguments***)
79 try:
---> 80 self._context.run(self._callback, *self._args)
81 except (SystemExit, KeyboardInterrupt):

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/kernelbase.py:510, in Kernel.dispatch_queue(***failed resolving arguments***)
509 try:
--> 510 await self.process_one()
511 except Exception:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/kernelbase.py:499, in Kernel.process_one(***failed resolving arguments***)
498 return None
--> 499 await dispatch(*args)

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/kernelbase.py:406, in Kernel.dispatch_shell(***failed resolving arguments***)
405 if inspect.isawaitable(result):
--> 406 await result
407 except Exception:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/kernelbase.py:730, in Kernel.execute_request(***failed resolving arguments***)
729 if inspect.isawaitable(reply_content):
--> 730 reply_content = await reply_content
732 # Flush output before sending the reply.

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/ipkernel.py:383, in IPythonKernel.do_execute(***failed resolving arguments***)
382 if with_cell_id:
--> 383 res = shell.run_cell(
384 code,
385 store_history=store_history,
386 silent=silent,
387 cell_id=cell_id,
388 )
389 else:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/ipykernel/zmqshell.py:528, in ZMQInteractiveShell.run_cell(***failed resolving arguments***)
527 self._last_traceback = None
--> 528 return super().run_cell(*args, **kwargs)

File ~/miniconda3/envs/default/lib/python3.10/site-packages/IPython/core/interactiveshell.py:2881, in InteractiveShell.run_cell(***failed resolving arguments***)
2880 try:
-> 2881 result = self._run_cell(
2882 raw_cell, store_history, silent, shell_futures, cell_id
2883 )
2884 finally:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/IPython/core/interactiveshell.py:2936, in InteractiveShell._run_cell(***failed resolving arguments***)
2935 try:
-> 2936 return runner(coro)
2937 except BaseException as e:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/IPython/core/async_helpers.py:129, in _pseudo_sync_runner(***failed resolving arguments***)
128 try:
--> 129 coro.send(None)
130 except StopIteration as exc:

File ~/miniconda3/envs/default/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3135, in InteractiveShell.run_cell_async(***failed resolving arguments***)
3133 interactivity = "none" if silent else self.ast_node_interactivity
-> 3135 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
3136 interactivity=interactivity, compiler=compiler, result=result)
3138 self.last_execution_succeeded = not has_raised

File ~/miniconda3/envs/default/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3338, in InteractiveShell.run_ast_nodes(***failed resolving arguments***)
3337 asy = compare(code)
-> 3338 if await self.run_code(code, result, async_=asy):
3339 return True

File ~/miniconda3/envs/default/lib/python3.10/site-packages/IPython/core/interactiveshell.py:3398, in InteractiveShell.run_code(***failed resolving arguments***)
3397 else:
-> 3398 exec(code_obj, self.user_global_ns, self.user_ns)
3399 finally:
3400 # Reset our crash handler in place

Input In [50], in ()
----> 1 jax.jacobian(ridge_solver)(10.)

File ~/miniconda3/envs/default/lib/python3.10/site-packages/jaxopt/_src/implicit_diff.py:251, in _custom_root..wrapped_solver_fun(***failed resolving arguments***)
250 keys, vals = list(kwargs.keys()), list(kwargs.values())
--> 251 return make_custom_vjp_solver_fun(solver_fun, keys)(*args, *vals)

JaxStackTraceBeforeTransformation: TypeError: missing a required argument: 'theta'

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

TypeError Traceback (most recent call last)
Input In [50], in ()
----> 1 jax.jacobian(ridge_solver)(10.)

File ~/miniconda3/envs/default/lib/python3.10/site-packages/jax/_src/api.py:1377, in jacrev..jacfun(*args, **kwargs)
1375 y, pullback, aux = _vjp(f_partial, *dyn_args, has_aux=True)
1376 tree_map(partial(_check_output_dtype_jacrev, holomorphic), y)
-> 1377 jac = vmap(pullback)(_std_basis(y))
1378 jac = jac[0] if isinstance(argnums, int) else jac
1379 example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args

[... skipping hidden 12 frame]

File ~/miniconda3/envs/default/lib/python3.10/site-packages/jaxopt/_src/implicit_diff.py:224, in _custom_root..make_custom_vjp_solver_fun..solver_fun_bwd(tup, cotangent)
221 else:
222 sol = res
--> 224 ba_args, ba_kwargs, map_back = _signature_bind_and_match(
225 reference_signature, *args, **kwargs)
226 if ba_kwargs:
227 raise TypeError(
228 "keyword arguments to solver_fun could not be resolved to "
229 "positional arguments based on the signature "
(...)
232 "custom_fixed_point if fixed_point_fun takes catch-all **kwargs, "
233 "both of which are currently unsupported.")

File ~/miniconda3/envs/default/lib/python3.10/site-packages/jaxopt/_src/implicit_diff.py:152, in _signature_bind_and_match(signature, *args, **kwargs)
150 args = [(False, i, v) for i, v in enumerate(args)]
151 kwargs = {k: (True, k, v) for (k, v) in kwargs.items()}
--> 152 ba = signature.bind(*args, **kwargs)
154 mapping = [(was_kwarg, ref) for was_kwarg, ref, _ in ba.args]
156 def map_back(out_args):

File ~/miniconda3/envs/default/lib/python3.10/inspect.py:3179, in Signature.bind(self, *args, **kwargs)
3174 def bind(self, /, *args, **kwargs):
3175 """Get a BoundArguments object, that maps the passed `args`
3176 and `kwargs` to the function's signature. Raises `TypeError`
3177 if the passed arguments can not be bound.
3178 """
-> 3179 return self._bind(args, kwargs)

File ~/miniconda3/envs/default/lib/python3.10/inspect.py:3094, in Signature._bind(self, args, kwargs, partial)
3092 msg = 'missing a required argument: {arg!r}'
3093 msg = msg.format(arg=param.name)
-> 3094 raise TypeError(msg) from None
3095 else:
3096 # We have a positional argument to process
3097 try:

TypeError: missing a required argument: 'theta'
```

This seems to conflict with what seems to be suggested in the paper.

Having played with the code a bit, it seems to be due to an unmatched number of arguments between the inner solver and the optimality condition `F`.

Let me note that if I test [this ridge regression example](https://github.com/google/jaxopt/blob/main/examples/implicit_diff/ridge_reg_implicit_diff.py), the `Optax` solver seems to be able to take gradient steps. **However**, I get unreasonable results for the Jacobian if I call the inner solver directly:
```
theta = 2.
data_tr = (data[0], data[2])
w_opt = ridge_solver(init_w, theta, data_tr)

# computing the Jacobian by calling F works fine
jax.jacobian(F)(w_opt, theta, data_tr)
Array([[2.00014853e+00, 4.33583318e-06, 1.17563381e-04, 2.49564437e-07,
4.56321777e-06, 4.07277803e-05, 6.18527411e-04, 1.38571322e-05,
1.41886936e-04, 4.17489372e-03, 1.33898677e-04, 1.96410366e-03,
1.34020054e-04],
[4.33583273e-06, 2.00294137e+00, 1.84755365e-04, 1.93276492e-06,
2.16197350e-05, 3.35783930e-04, 1.81916286e-03, 3.17757367e-04,
1.92739695e-04, 1.40373148e-02, 8.29683442e-04, 1.92523114e-02,
3.15814861e-04],
[1.17563381e-04, 1.84755365e-04, 2.00046730e+00, 2.56317730e-06,
1.93931883e-05, 2.03579679e-04, 2.49244971e-03, 1.04629522e-04,
3.09628609e-04, 1.35955038e-02, 6.11885451e-04, 1.14722112e-02,
4.62129770e-04],
[2.49564437e-07, 1.93276492e-06, 2.56317730e-06, 2.00000024e+00,
1.29375820e-07, 1.55074326e-06, 1.70912572e-05, 8.09174367e-07,
1.77757283e-06, 7.96649838e-05, 4.17608953e-06, 9.16967983e-05,
2.64862501e-06],
[4.56321777e-06, 2.16197368e-05, 1.93931883e-05, 1.29375834e-07,
2.00000095e+00, 1.17004456e-05, 1.26681291e-04, 6.94196069e-06,
1.43129946e-05, 6.77908247e-04, 3.34996876e-05, 6.66017295e-04,
2.22999697e-05],
[4.07277803e-05, 3.35783930e-04, 2.03579679e-04, 1.55074326e-06,
1.17004456e-05, 2.00014234e+00, 1.40636507e-03, 9.05015477e-05,
1.51485103e-04, 7.61069637e-03, 3.97522817e-04, 8.12703930e-03,
2.39455359e-04],
[6.18527411e-04, 1.81916275e-03, 2.49244971e-03, 1.70912572e-05,
1.26681291e-04, 1.40636507e-03, 2.01715755e+00, 7.60396302e-04,
1.80630060e-03, 8.37670639e-02, 4.06576833e-03, 7.98479617e-02,
2.90394668e-03],
[1.38571322e-05, 3.17757367e-04, 1.04629522e-04, 8.09174367e-07,
6.94196069e-06, 9.05015622e-05, 7.60396360e-04, 2.00007224e+00,
7.97978137e-05, 4.49323049e-03, 2.54728773e-04, 5.30747883e-03,
1.38489166e-04],
[1.41886936e-04, 1.92739681e-04, 3.09628609e-04, 1.77757272e-06,
1.43129946e-05, 1.51485088e-04, 1.80630060e-03, 7.97978137e-05,
2.00033450e+00, 1.14652235e-02, 4.56731883e-04, 8.11540522e-03,
3.35147721e-04],
[4.17489372e-03, 1.40373148e-02, 1.35955038e-02, 7.96649838e-05,
6.77908247e-04, 7.61069637e-03, 8.37670565e-02, 4.49323002e-03,
1.14652226e-02, 2.49225712e+00, 2.22632252e-02, 4.27025557e-01,
1.51680810e-02],
[1.33898677e-04, 8.29683500e-04, 6.11885451e-04, 4.17608999e-06,
3.34996876e-05, 3.97522817e-04, 4.06576833e-03, 2.54728773e-04,
4.56731883e-04, 2.22632233e-02, 2.00115418e+00, 2.29516216e-02,
7.26781436e-04],
[1.96410366e-03, 1.92523114e-02, 1.14722103e-02, 9.16967983e-05,
6.66017295e-04, 8.12703930e-03, 7.98479617e-02, 5.30747883e-03,
8.11540522e-03, 4.27025557e-01, 2.29516216e-02, 2.48472643e+00,
1.37193939e-02],
[1.34020011e-04, 3.15814774e-04, 4.62129887e-04, 2.64862501e-06,
2.22999697e-05, 2.39455316e-04, 2.90394644e-03, 1.38489180e-04,
3.35147721e-04, 1.51680792e-02, 7.26781203e-04, 1.37193957e-02,
2.00059700e+00]], dtype=float32)

# calling inner solver to compute jacobian returns 0, regardless of the input parameters
jax.jacobian(ridge_solver)(init_w, theta, data_tr)
Array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]], dtype=float32)
```

Can anyone help me understand what I am doing wrong?

Thank you very much!