Scan trick to improve compilation time
patrick-kidger opened this issue · comments
It should be doable to apply equinox.internal.scan_trick inside CG and BiCGStab.
patrick-kidger / lineax
Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax