Enable setting of Jacobian tags in ODETerm

[jpbrodrick89]

In light of the new colouring methods in lineax I really wanted to see if there was an unfootgunnable way for users to provide structural tags for their syste, and I think I might have found it.

The challenge

Providing tags to implicit solvers requires user's knowledge of the exact impicit relation being solved and perhaps careful consideration of whether the solve is fully-implicit (in which case a tridiagonal ODE term may be become pentadiagonal for a FIRK2) or direct implicit.Fortunately though, diffrax currently only has DIRK solvers (Implicit Euler is just a flavour of DIRK).

The observation

Each iteration of the nonlinear solver solves a linearised form of the implicit relation—this can be represented in terms of lineax operators and our existing tag propagation logic.

The solution

Users provide tags for the ODE system Jacobian directly on to ODETerm, this should represent the structure of jacobian(vf)(y) and is completely agnostic of solver internals. Implicit solvers use lineax tag propagation to determine which tags remain relevant. For example, if a user declares a system to be negative semi-definite (e.g. a diffusion equation) then the current DIRK and implicit Euler solvers will recognised this as positive/negative semi-definite and pick Cholesky, however if the system is positive semi-definite (anti-diffusion) the implicit solvers will realise that residual jacobian is not guaranteed to have either form of semi-definiteness for arbitrary timestep and LU will be used instead.

Remaining footguns

The inherent footgun is not necessary completely removed, but the burden of complexity now lays solely at the feet of power users developing their own implicit solvers rather than the typical diffeqsolve who can just plug in well designed implicit solvers and have it "just work". The main responsibility of implicit solver writers is to provide implicit_op_relation mirroring implicit_relation. Furthermore, this handling is opt-in from implicit solver designers, no mandatory (or optional) arguments are introduced into implicit solver design.

It would be really nice if we can get this in before the AD workshop and provide an easy heat equation example but completely understand if this is too complicated to find the time for in that window.