Including user-defined Jacobian

[Justin-Tan]

Hi devs, looks like a really nice library. I've been looking for a Jax-native root finding method that supports vmap for some time. Currently I am using an external call to scipy.optimize.root together with the multiprocessing library, which is quite slow.

The runtime for root finding using the Newton method in this library is slower than the above method though - I suspect this is because the Jacobian needs to be calculated at each iteration. Is there a way for the user to supply an analytic Jacobian? Or could you point me in the right direction to implement this feature?

For reference, this is my MWE in case I am not doing things efficiently:

from jax import jit, jacfwd, vmap, random
import optimistix as optx

def fn(y, b):
    return (y-b)**2

M = 1024
key = random.PRNGKey(42)
key, key_ = random.split(key, 2)

y = random.normal(key, (M,))
b = random.normal(key_, (M,))
sol = optx.root_find(vmap(fn), solver, y, b)

[patrick-kidger]

Okay, many things to respond to here!

Speed

With respect to the speed, for your JAX code are you:

• JIT'ing everything;
• excluding compile time;
• including block_until_ready?

In practice this means writing things out something like:

@jax.jit
def run(y, b):
    sol = optax.root_find(vmap(fn), solver, y, b)
    return sol.value

run(y, b)  # compile
times = timeit.repeat(lambda: jax.block_until_ready(run(y, b)), number=1, repeat=10)
print(min(times))

Recalculating Jacobians

You commented on calculating the Jacobian afresh every iteration. If using the typical Newton algorithm then this is expected (desired) behaviour. But if you're saying that you'd prefer to use a quasi-Newton algorithm like the chord method (that computes the Jacobian once at the initial point and then re-uses it), then there is optx.Chord as well.

Analytical Jacobians

You commented on supplying an analytical Jacobian. This isn't necessary, as the analytical Jacobian is actually already derived from fn automatically using autodifferentiation. Unless the autodiff does something surprisingly inefficient, then providing one manually wouldn't meaningfully improve things there.

Custom Jacobians

If despite everything you really do want to provide a custom Jacobian, then this can be done using jax.custom_jvp. By wrapping your fn in a jax.custom_vjp, then you can override how JAX calculates autoderivatives of your code. (And this will then be picked up by the autodiff used by Optimistix to calculate the Jacobian.)

Does the above help?

[itk22]

Hi @patrick-kidger ,
I might have found a case where supplying an analytical Jacobian is crucial for good performance. In my application, I have to deal with large Jacobians, which need to be stored as sparse arrays to prevent OOMs. Even though I have a formula for calculating the sparse Jacobians, including it in the custom_jvp will lead to a dense output as I described here.

[jpbrodrick89]

One way to do this would be to create a custom linear solver that can deal with the sparse representation of your matrix. For example, see my Lineax PR which uses coloring methods to efficiently compute and solve a tridiagonal system without ever instantiating the full jacobian.

However, if you need a dense solver this may not be possible without a lot of work.

(I plan to make the request changes to the above PR this or next week but have been delayed due to competing priorities.)

[johannahaffner]

There are two separate concerns here - preserving sparsity when computing the Jacobian, and then using a linear solver that preserves said sparsity and does not materialise a dense system, assuming that you are interested in using optx.Newton to do the actual root-finding operation.