Split coloring and jacrev/jacfwd

[nardi]

Hi, I wanted to use this library but for efficiency purposes I wanted to split the graph coloring step from the actual evaluation of the Jacobian.

This is because I am calculating many Jacobians with the same sparsity pattern, and also doing a nested differentiation, e.g. using jacfwd to calculate the sparse Jacobian matrix, solving a linear system involving it using a sparse linear solver, and then using jax.jvp to differentiate through that solve. This means the jacfwd(f, sparsity=...) call is inside a jax.jit context, which means the sparsity pattern array needs to be passed through as a static argument, which is not really nice (since arrays are not hashable).

I had a go at adding a feature to support this. I added ForwardModeColoring and ReverseModeColoring objects, and functions to produce them. These are PyTrees (easy to JIT) and can be passed to jacfwd/rev as jacfwd(f, coloring=...). I also made the representation quite compact, in that e.g. the projection matrix for the reverse mode coloring is not fully constructed. This is quite a clean way to split the coloring and the evaluation, IMO :)

I used LLM as a coding assistant, but reviewed all code myself, and it of course passes all existing tests and the new ones added to test the split functionality specifically. I also added tests specific to my use case, for example test_jacfwd_with_precomputed_coloring_jit, where the coloring is computed outside of JIT and then passed into a JIT-compiled function.