search

patrick-kidger / optimistix

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/
Apache License 2.0
284 stars 14 forks source link

Add transpose rule for stop_gradient #68

Closed dkweiss31 closed 1 month ago

dkweiss31 commented 1 month ago

Work around for JAX issue #22011, as well as https://github.com/patrick-kidger/diffrax/pull/387#issuecomment-2174488365

patrick-kidger commented 1 month ago

So in direct contradiction to my past self...

https://github.com/patrick-kidger/diffrax/pull/387#issuecomment-2199206730

... I think this should actually be in Lineax!

Sorry, I don't know why I said Optimistix before. It's Lineax that does the linearisation/transposition! 😅

Would you be able to re-open this somewhere in Lineax instead? (FWIW I can at least give you a review now: LGTM! :D)

dkweiss31 commented 1 month ago

Yes of course!