Internal: Delete unused code

docs/examples_parameter_estimation/neural_ode.py

@@ -91,7 +91,7 @@ def build_update_fn(*, optimizer, loss_fn):
     @jax.jit
     def update(params, opt_state):
         """Update the optimiser state."""
-        loss, grads = jax.value_and_grad(loss_fn)(params)
+        _loss, grads = jax.value_and_grad(loss_fn)(params)
         updates, opt_state = optimizer.update(grads, opt_state)
         params = optax.apply_updates(params, updates)
         return params, opt_state

docs/examples_parameter_estimation/physics_enhanced_regression_1.py

@@ -142,7 +142,7 @@ def build_update_fn(*, optimizer, loss_fn):
     @jax.jit
     def update(params, opt_state):
         """Update the optimiser state."""
-        loss, grads = jax.value_and_grad(loss_fn)(params)
+        _loss, grads = jax.value_and_grad(loss_fn)(params)
         updates, opt_state = optimizer.update(grads, opt_state)
         params = optax.apply_updates(params, updates)
         return params, opt_state

probdiffeq/backend/containers.py

@@ -9,8 +9,3 @@
 @dataclass_transform()
 def dataclass(*args, **kwargs):
     return dataclasses.dataclass(*args, **kwargs)
-
-
-@dataclass_transform()
-def dataclass_astuple(datacls):
-    return dataclasses.astuple(datacls)

probdiffeq/backend/linalg.py

@@ -54,10 +54,6 @@ def vector_norm(arr, /, *, order=None):
     return jnp.linalg.norm(arr, ord=order)
 
 
-def matrix_norm(arr, /, *, order=None):
-    return jnp.linalg.norm(arr, ord=order)
-
-
 def solve_triangular(matrix, rhs, /, *, trans=0, lower=False):
     return jax.scipy.linalg.solve_triangular(matrix, rhs, trans=trans, lower=lower)
 

probdiffeq/impl/_conditional.py

@@ -2,16 +2,11 @@
 
 from probdiffeq.backend import abc, containers, functools, linalg, tree_util
 from probdiffeq.backend import numpy as np
-from probdiffeq.backend.typing import Any, Array, Callable
+from probdiffeq.backend.typing import Any, Array
 from probdiffeq.impl import _normal
 from probdiffeq.util import cholesky_util, linop_util
 
 
-class Transformation(containers.NamedTuple):
-    matmul: Callable
-    bias: Array
-
-
 class TransformBackend(abc.ABC):
     @abc.abstractmethod
     def marginalise(self, rv, transformation, /):
@@ -165,7 +160,7 @@ def conditional(self, matmul, noise):
         return Conditional(matmul, noise)
 
     @abc.abstractmethod
-    def identity(self, num_derivatives_per_ode_dimension, /):
+    def identity(self, ndim, /):
         raise NotImplementedError
 
     @abc.abstractmethod

probdiffeq/impl/_normal.py

@@ -18,7 +18,7 @@ def preconditioner_apply(self, rv: Normal, p, /) -> Normal:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def standard(self, num_derivatives_per_ode_dimension, /, output_scale) -> Normal:
+    def standard(self, num, /, output_scale) -> Normal:
         raise NotImplementedError
 
 

probdiffeq/util/cholesky_util.py

@@ -106,20 +106,6 @@ def revert_conditional(R_X_F, R_X, R_YX):
     return R_Y, (R_XY, G)
 
 
-def _triu_via_shortcut(R, R_YX):
-    """If R_X and R_X_F are zero, triu_via_qr(R) can be implemented cheaply.
-
-    Namely, by applying it only to R_YX and embedding the result in zeros.
-    This is not only more efficient because it requires fewer floating-point operations
-    than qr-decomposing the full matrix, but it also admits a well-defined
-    reverse-mode derivative!
-    """
-    R = np.zeros_like(R)
-    R_YX = triu_via_qr(R_YX)
-    n, m = np.shape(R_YX)
-    return R.at[:n, :m].set(R_YX)
-
-
 def _is_matrix(mat, matrix_ndim=2):
     return np.ndim(mat) == matrix_ndim