Differentiating a solver for roots of a complex polynomial

[fbartolic]

Hi all! I'm trying to implement an astrophysics model in JAX. The central part of the model involves solving for the roots of a 5th order complex polynomial and I need to do that for a large grid of complex numbers in parallel on a GPU. jax.numpy.roots works but it relies on np.linalg.eig which isn't supported on the GPU (see #1259). Is there an alternative algorithm which would work on a GPU?

I tried implementing a simple solver using Laguerre's method in JAX but I'm not sure how to implement the tangent_solve method required by lax.custom_root, also, my naive implementation isn't jit-able so it's orders of magnitude slower than jax.numpy.roots.

Any help would be much appreciated :).

[fbartolic]

I managed to write a JIT-able version of a solver based on Laguerre's method by replacing the if statements with lax.cond and the for loop with a lax.while_loop. It works on a GPU and there's still plenty of room for improvement in terms of speed.

I've done some reading on propagating gradients through a solver and if I understand it correctly, the implicit function theorem tells me that to compute the gradient of the roots z^* with respect to model parameters (polynomial coefficients), I only need to evaluate the first derivative of my complex polynomial p(z) at the roots z^* which is just a 4th order complex polynomial. I can do this by either writing a custom vjp function as in this tutorial or I can use the more convenient lax.custom_root function which does that for me.

I'm still confused about how to use lax.custom_root. It requires a root solver function of the form solve(f, initial_guess) where f in my case is the complex polynomial, however my root solver complex_roots(p) doesn't work by evaluating the polynomial function directly but rather it takes the coefficients p of the polynomial and then does a bunch of things with those coefficients (it's an old numerical recipes routine). How can I adapt this to a form I could pass to lax.custom_root? Also, I don't really understand the connection between tangent_solve and the Jacobian of the root solutions that evaluated at the roots.

[shoyer]

You don't need to make use of f an initial_guess inside the solve function passed to custom_root -- it's OK to pass in an arbitrary callable that computes the root based on closed over variables like the vector of coefficients p.

What tangent_solve needs to do is take a function that applies the Jacobian J, and then solve a matrix-free linear system of equations J(x) = b. Depending on the dimensionality of your polynomials, it may be just fine to compute a dense Jacobian with jax.jacobian and use jax.numpy.linalg.solve.

This example from the JAX test suite may be a useful reference point:

    import jax.numpy as jnp
    import jax
    from jax import lax

    def vector_solve(f, y):
      return jnp.linalg.solve(jax.jacobian(f)(y), y)

    def linear_solve(a, b):
      f = lambda y: high_precision_dot(a, y) - b
      x0 = jnp.zeros_like(b)
      solution = jnp.linalg.solve(a, b)
      oracle = lambda func, x0: solution
      return lax.custom_root(f, x0, oracle, vector_solve)

linear_solve(a, b) is equivalent to jnp.linalg.solve(a, b) with gradients are calculated via custom_root. The function f here happens to be affine, but it could be replaced any arbitrary alternative.

[fbartolic]

Thank you! I did manage to implement a custom jvp rule which worked in the end but I'll try the custom_root approach as well to compare.