Unify all 'conditional' types (#832)

* Reorganise some code to improve readability

* Move cholesky_hilbert to cholesky_util and rename Conditional -> LatentCond to prepare refactor

* Update large chunks of Dense and Isotropic conditionals

* Update blockdiag transitions

* Conditionals are LatentCond's now (consistently)

* Delete strategy.extrapolate_mean because it's the same for all strategies now

* Remove a useless 'name' argument

* Improve some docstrings

* Delete commented-out code and revert tutorial to previous state

Author: pnkraemer

docs/examples_basic/conditioning_on_zero_residual.py

@@ -56,7 +56,7 @@ def vector_field(y, t):  # noqa: ARG001
 tcoeffs_like = [u0] * (NUM_DERIVATIVES + 1)
 ts = jnp.linspace(t0, t1, num=500, endpoint=True)
 init_raw, transitions, ssm = ivpsolvers.prior_wiener_integrated_discrete(
-    ts, tcoeffs_like=tcoeffs_like, output_scale=100.0, ssm_fact="dense"
+    ts, tcoeffs_like, output_scale=100.0, ssm_fact="dense"
 )
 
 markov_seq_prior = stats.MarkovSeq(init_raw, transitions)

probdiffeq/backend/linalg.py

@@ -67,5 +67,9 @@ def diagonal(arr, /):
     return jnp.diagonal(arr)
 
 
+def diagonal_matrix(arr, /):
+    return jnp.diag(arr)
+
+
 def triu(arr, /):
     return jnp.triu(arr)

probdiffeq/backend/tree_util.py

@@ -28,5 +28,9 @@ def ravel_pytree(tree, /):
     return jax.flatten_util.ravel_pytree(tree)
 
 
+def tree_flatten(tree, /):
+    return jax.tree_util.tree_flatten(tree)
+
+
 def register_dataclass(datacls):
     return jax.tree_util.register_dataclass(datacls)