search

patrick-kidger / lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax
Apache License 2.0
365 stars 24 forks source link

Follow conjugation convention in `tree_dot` #105

Closed Randl closed 3 months ago

Randl commented 3 months ago

It may break something downstream unfortunately, so updating dependencies should be done carefully

patrick-kidger commented 3 months ago

Modulo failing tests -- not sure what's going on with those (?) -- this LGTM. Once the tests are fixed we can check the Optimistix and Diffrax tests on this PR before merging it, and see if anything looks awry.

Randl commented 3 months ago

Should we bump the version since this is a potentially breaking change?

patrick-kidger commented 3 months ago

Yes, I think we should.

I can see that tests now pass -- are you happy with this PR / should I merge it?

Randl commented 3 months ago

I think I am. Though one possibility to consider is to not do conjugation in tree_dot at all but rater leave it to caller, similarly to jnp.dot. This may clutter downstream code but will be consistent with the Jax/numpy semantics.

patrick-kidger commented 3 months ago

Merged, then! And that's an interesting idea. In particular given the difficulties we faced in https://github.com/patrick-kidger/optimistix/pull/61, I can see that that might simplify things.

The decision is yours -- I'd be happy to make that change if you think that is what is best.