3D heat diffusion

Please excuse the somewhat open-ended question!

Could you comment on the suitability of Diffrax for modelling the transient heat diffusion equation in 3D? I have followed the 1D example in the docs, where you obtain a system of coupled ODEs by discretising the spatial domain. I understand this could be extended to 3D space, indeed Crank-Nicolson is often applied to 3D space, but I am aware it would require me to write my own solver. I wonder whether this is not the sort of thing you originally wrote Diffrax for. Other libraries (FEniCSx, FiPy) exist for this sort of thing, but they are outside the JAX ecosystem. Diffrax still appeals for performance reasons and because I like the documentation.

Start with the linked nonlinear_heat_pde documentation example and compare its 1D spatial discretisation with the requested 3D transient heat-diffusion use case. Define whether supporting this belongs in Diffrax and what a completed contribution would need; the issue does not identify files, tests, or a concrete implementation scope.