lax.root, a primitive for differentiable root finding

[shoyer]

This should solve the issue with closed over variables not being handled properly in _custom_implicit_solve.

Does implementation look sane? It's a little messier than I would like, because we need to deal with inputs into provided functions.

See the tests for an example of what using this API looks like. There are lots of ways to make this more user-friendly, but this is intended as a low-level API for defining implicit derivatives. The user facing APIs will be routines like scipy.optimize.root, which won't require providing functions for solve and tangent_solve.

TODOs:

• replace asserts with errors
• more test coverage
  • something checking higher dimensional arrays / nested structures
  • check grads with jtu.check_grads
  • verify the jit works properly

[shoyer]

The need for this suggests that I'm doing something wrong...

[mattjj]

Zeros shouldn't escape the AD system, i.e. they should only appear in JVP rules. Maybe a check for ad_util.zero just needs to be added upstream, i.e. in the caller (assuming this is called from a JVP rule).

[shoyer]

I tried using ad.linearize, but the jaxpr it returns isn't a TypedJaxpr, so I couldn't figure out how to evaluate it.

[mattjj]

I think we might want to use ad.jvp_jaxpr here, basically following the logic of scan, though I'm not sure yet if we need the fixed-point logic. I can explain more about that in chat.

[mattjj]

Nevermind about that fixed-point stuff; that's not relevant, because the variables we're differentiating with respect to are all in the closure of the function being passed in.

[mattjj]

I'm paging this stuff back in, and more exited about it than ever.

Here are some things I want to check with you: my current best understanding is we want a primitive

root :: (a -> a) -> a -> a

where the first argument is a function which we can map to math as f : R^n -> R^n, the second argument is an initial guess x_0 (a point in R^n), and the output x* is a point in R^n satisfying f(x*) = 0 (where the RHS is the zero vector in R^n).

We expect to apply root to functions that have nontrivial closures, and in particular that might close over values involved in differentiation. So while the above is a description at the API level, really we should think of the mathematical function in a closure-converted way as f : R^n x R^p -> R^n, where the first argument is some set of parameters. Then we can say x*(a) solves f(x*(a), a) = 0.secondI think we want the JVP rule to look something like

root_jvp f x0 a adot =
  let x_star = root (lambda x: f(x, a)) x0
      x_star_dot = linear_solve (∂_0 f (x_star, a)) (∂_1 f(x_star, a)[-a_dot])
  in (x_star, x_star_dot)

where I'm using square brackets to denote the application of the linear function ∂_1 f(x_star, a) to a vector.

Does that sound right?

[shoyer]

Yes, this looks right to me, with one minor correction:

So while the above is a description at the API level, really we should think of the mathematical function in a closure-converted way as f : R^n x R^p -> R^n, where the firstsecond argument is some set of parameters

[mattjj]

LGTM! Thanks for pushing on this. We've got some more polishing work to do, but this is a big step forward: it actually works!

(We've been discussing extensively offline what follow-up stuff we want to do.)

[mattjj]

I think this check is redundant because you already checked it when you formed jaxpr. It doesn't hurt to include though, other than taking up precious vertical space :)

[shoyer]

We evaluated f() when we formed the jaxpr, but not solve(). So I think we do need this. Actually I even wrote a test for that catches this error message :)

[mattjj]

We think we can avoid instantiating zeros here, and otherwise be more conservative about how much work we do (I currently think we do want to run a fixed-point on ad.jvp_jaxpr!), but we can leave that for follow-up work. Land and iterate!