Scalar and matrix-valued differential equations

[pnkraemer]

What?

Would be great to be able to support equations of the form M'(t) = f(M(t)) where $M$ is an arbitrarily-shaped array (e.g., a matrix) and not necessarily a (n,)-shaped array.

Why?

This is super useful when solving for instance partial differential equations and has come up recently in #452.

Context

While we can always "flatten" the equation, it is not just easier to use the code if flattening isn't necessary, but matrix-valued equations also motivate a different state-space model and maybe we can improve efficiency here.

Other

Stretch goal: solve equations where the state is an arbitrary pytree. It is not obvious, however, what the state-space model should be for an arbitrary pytree so this stretch goal is likely out of scope.

Related

If #288 can be realised successfully, we get matrix-valued differential equations for free. See also #296, which "goes the other way" (i.e. demanding solvers for ndim-0 arrays instead of ndim-1 arrays).