search

patrick-kidger / optimistix

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

Usage with 'vmap' #16

Closed Justin-Tan closed 9 months ago

Justin-Tan commented 9 months ago

Hi, looks like a very promising library, this bit in the docs got me interested:

Unlike the SciPy implementation of Newton's method, the Optimistix version also works for vector-valued (or PyTree-valued) y.

Does this mean that the function passed to the root finder has to be vectorized in the traditional sense? Do the root finders here support functions which rely on vmap? I couldn't find anything in the docs about this. Thanks!

Justin-Tan commented 9 months ago

Nevermind, doing the obvious thing and wrapping function with vmap seems to work.