Implement support for pytree-valued differential equations (#833)

* Taylor series estimation supports arbitrary pytrees now

* Run Taylor tests for all state-space model factorisations

* Start making the solvers pytree-compatible

* I need a break...

* Only left to update the stats (and the other factorisations)

* Update stats

* Fix terminal-value-LML tests

* Update LML code

* Update remaining stats modules

* Remove print statements from the src

* Fix the last dense test

* Update isotropic implementations

* All tests pass for isotropic & dense implementations

* Update blockdiagonal implementation

* Rerun benchmarks

* Change taylor-coeff notebook to logistic ODE to demonstrate scalar-valued problems

* Delete outdated comments

* Clean up comments

* Delete dead code

* Improve formatting

Author: pnkraemer

docs/benchmarks/hires/results.npy

@@ -40,7 +40,7 @@
 # We start by defining an ODE.
 
 # +
-f, u0, (t0, t1), f_args = ivps.rigid_body()
+f, u0, (t0, t1), f_args = ivps.logistic()
 
 
 def vf(*y, t):  # noqa: ARG001

docs/benchmarks/lotkavolterra/results.npy

@@ -2,7 +2,7 @@ format-and-lint:
 	pre-commit run --all-files
 
 test:
-	pytest -n auto -v -Werror 		# parallelise, verbose output, warnings as errors
+	pytest -n auto -v -Werror		# parallelise, verbose output, warnings as errors
 
 quickstart:
 	# Run some code without installing any of the optional dependencies

docs/benchmarks/pleiades/results.npy

@@ -133,11 +133,11 @@ def eye(n, m=None, /):
 
 
 def save(path, arr, /):
-    return jnp.save(path, arr)
+    return jnp.save(path, arr, allow_pickle=True)
 
 
 def load(path, /):
-    return jnp.load(path)
+    return jnp.load(path, allow_pickle=True)
 
 
 def allclose(a, b, *, atol=1e-8, rtol=1e-5):