Matrix inverse of pytree of matrices

[f0uriest]

I'm trying to implement some Newton type methods for pytrees. I know there are several implementations that use iterative linear algebra (like cg or gmres) to solve for the Newton step, but I'm trying to get it to work with direct methods which are faster and more accurate for my applications.

Basically, I have a function f whose argument is a pytree and outputs a scalar. I can get the Hessian of f which is then a pytree of pytrees of matrices. How can I invert this matrix or do something like jnp.linalg.solve?

Relatedly would be doing least squares with the Jacobian of a vector valued function using jnp.linalg.qr or jnp.linalg.svd

[mattjj]

Thanks for the question!

Would you be willing to concatenate all the blocks together into a big 2D array, apply direct methods there, and then take it back apart into the pytree-of-pytrees-of-blocks representation? If so, you could use the jax.flatten_util.ravel_pytree function.

If you want to keep things organized as pytrees-of-pytrees-of-blocks, then JAX doesn't offer any functions for operating on those matrix representations. You'd have to write the direct method algorithm yourself, e.g. operating on lists-of-lists-of-Arrays, though you could call jax.numpy.linalg functions on the Arrays as subroutines. (@shoyer had an experiment in this direction at one point, maybe in a PR or issue, but I don't think it could handle this case.)

What do you think?

[f0uriest]

concatenating everything into a big 2d array would certainly work as a first pass. It seems like ravel_pytree would flatten everything down to 1d though? Is there a way to flatten a pytree down to 1d, but a nested pytree down to 2d in a self consistent way?

[patrick-kidger]

Try Lineax, which in particular offers way to represent pytree linear operators.

For what it's worth, for your particular problem, this is mostly a convenience: with a direct method (e.g. an LU solver), then it still calls out to JAX's underlying routines, doing a ravel_pytree under the hood. (With the appropriate details to produce something in 2D rather than 1D, as you touch on in your last message.)

[f0uriest]

wow, Lineax seems really cool! What would be the best way to do a Hessian of a scalar function? As something like PyTreeLinearOperator(jax.hessian(f)(x)) or as JacobianLinearOperator(jax.grad(f), x)?

[patrick-kidger]

Thank you!
As for your question, both approaches will work! The best choice depends whether you want to materialise the Hessian (the first option), or leave it unmaterialised (the second option). Which one to use represents a trade-off between runtime/compiletime/memory-usage.

FWIW we should probably add a HessianLinearOperator(f, x), that does the second thing. (And lineax.materialise would convert it to a PyTreeLinearOperator.) That's currently tracked in patrick-kidger/lineax#4.