Complex type support #326
Complex type support #326
Conversation
| return jsp.linalg.lu_factor(jax.jacfwd(curried)(flat)) | ||
| return jsp.linalg.lu_factor( | ||
| jax.jacfwd(curried, holomorphic=jnp.iscomplexobj(flat))(flat) | ||
| ) |
There was a problem hiding this comment.
Given the guard against implicit solvers earlier, is this change necessary?
(I think we'd want a lot of tests to be sure that the nonlinear solvers handle complex numbers correctly, what with all the linear solves and such going on.)
There was a problem hiding this comment.
With this change, implicit solvers are working for me locally. If anything, I'd remove the guard?
There was a problem hiding this comment.
So as it happens, I'm currently in the process of upgrading Diffrax to use Lineax and Optimistix to handle the linear and nonlinear solving. This is going to be great from a feature point of view, but it does also mean that any changes we make here will be short-lived.
Other than that, I think we'd need some fairly comprehensive tests to check that things are working. There's subtleties here that can easily go wrong, e.g. transpose vs conjugate-transpose. (JAX does the former, PyTorch does the latter.)
To justify the above, btw -- the position I've been taking with the JAX sciML ecosystem is that it's better not to do something than it is to do it unreliably. Hence why complex support has been such a slow-moving thing to have happen!
There was a problem hiding this comment.
I agree with you regarding taking slow part. But this code otherwise will just crash, and it passes basic tests. If you insist I'll remove it, or reinstate the guard
There was a problem hiding this comment.
Yup, I'd reinstate the guard for now. (As I say, the current implementation will be deleted in ~2 weeks anyway.)
I think getting complex support into Lineax will be the main thing needed to get proper support here.
|
Closing in favour of #330, which fixes the pre-commit failure present here. Thank you to all of @TimonHoess @Randl @andyElking for your efforts getting this in! |
Rebase of #197