[WIP] minimization with complex arguments

[Randl]

Following up on #61 (comment)
This doesn't introduce fixes yet, just adds a simple test that fails to highlight the issue

[Randl]

Locally, with patrick-kidger/lineax#103 this fails three least square tests (two are related to some pytree structure, and one should be some complex-related bug since the wrong answer is returned) and most of the minimize_jvp tests.

[Randl]

Unfortunately, jvp is frustratingly hard to debug; for example, the call stack is useless due to all the wrappers around the function calls. @patrick-kidger, do you have any advice on how I can approach that?

Specifically, I see that the LU solver gets 0 instead of 2 as a matrix, but I have no idea where it is called from and what this value is.

We can, of course, go the old way and merge the working stuff (minimization as for now and hopefully least squares soon) and work out the jvp later.

[patrick-kidger]

So the stack goes function -> jaxpr -> runtime. I think if the trace-time callstack doesn't help you, then I do have some tricks up my sleeve to help with the jaxpr and runtime approaches.

To help with intercepting things at the jaxpr level then I like to use eqx.debug.announce_transform. Passing an appropriate announce function then you can e.g. insert a breakpoint into the point at which a JVP transform is performed. (Even if this happens after a function is parsed into a jaxpr.)

As for the runtime level then lots of jax.debug.prints / jax.debug.breakpoints would probably be my usual thing here -- bisect through the program.

[Randl]

Ok, it turns out many tests passed since I initialized the minimization at the real line. I've changed that, and now there are more fails, but that helped me to find where the problem is.