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

Moved state-transpose to avoid recomputation #47

Closed packquickly closed 1 year ago

packquickly commented 1 year ago

Noticed the computation of solver.transpose(state, options) was duplicated in the dependent columns + dependent rows case. Moved it to avoid this duplication.

patrick-kidger commented 1 year ago

LGTM! The test failures look spurious -- due to a bug in upstream JAX, already fixed in https://github.com/google/jax/pull/17684. (+the fact that we depend on an unreleased version of Equinox).