has_aux option in custom_root
Open
- Dominant language
- Python
- Stars
- 1.1k
- Forks
- 76
- Avg merge
- 2d 21h
- Merged PRs (30d)
- 1
Description
Hi!
I'm trying to solve a bilevel problem involving a neural network using implicit diff, but I have an incompatible shape in `solve_normal_cg` since the matrix multiplying in `matvec` is not square.
As explained in https://github.com/google/jaxopt/issues/183, this can be solved by giving an initialization to the solver with the correct shape, but I have no clue how to find the correct shape.
Is there an automatic way to infer the shape of the initialization?
Thanks!
Contributor guide
Assessment
This issue has not been assessed yet.