Outsource Jacobian handling from linearisations (#847)

* Move ssm._linearize content to ssm._conditional because they operate on the same datastructures

* Implement Jacobian handling in a dedicated object to clarify stochastic trace estimation

* Separate Hutchinson from materialisation and forward/reverse autodiff

* Make all JacobianHandler arguments keyword-args for now

Author: pnkraemer

probdiffeq/backend/func.py

@@ -28,6 +28,10 @@ def linearize(func, *args):
     return jax.linearize(func, *args)
 
 
+def vjp(func, *args):
+    return jax.vjp(func, *args)
+
+
 def jvp(func, /, primals, tangents):
     return jax.jvp(func, primals, tangents)
 

probdiffeq/backend/typing.py

@@ -1,7 +1,7 @@
 """Typing module."""
 
 from collections.abc import Callable, Sequence
-from typing import Any, Generic, Protocol, TypeVar
+from typing import Any, Generic, Literal, Protocol, TypeVar
 
 from jax import Array
 from jax.typing import ArrayLike