Run all tests in single precision (#835)

* Move allclose() from backend.numpy to backend.testing

* There is only a single testing.allclose function now

* All tests run in single precision now

* Run all examples in single precision

Author: pnkraemer

docs/examples_advanced/equinox_while_loop.py

@@ -27,9 +27,6 @@
 from probdiffeq import ivpsolve, ivpsolvers, taylor
 from probdiffeq.backend import control_flow
 
-jax.config.update("jax_platform_name", "cpu")
-
-
 # -
 
 # Overwrite the while-loop (via a context manager):

docs/examples_advanced/parameter_estimation_blackjax.py

@@ -135,11 +135,6 @@
 from probdiffeq import ivpsolve, ivpsolvers, stats, taylor
 
 # +
-# x64 precision
-jax.config.update("jax_enable_x64", True)
-
-# CPU
-jax.config.update("jax_platform_name", "cpu")
 
 # IVP examples in JAX
 if not backend.has_been_selected:

docs/examples_advanced/parameter_estimation_optax.py

@@ -37,8 +37,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_enable_x64", True)
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 

docs/examples_basic/conditioning_on_zero_residual.py

@@ -32,9 +32,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_platform_name", "cpu")
-jax.config.update("jax_enable_x64", True)
-
 
 # +
 # Create an ODE problem

docs/examples_basic/dynamic_output_scales.py

@@ -46,7 +46,6 @@
     backend.select("jax")  # ivp examples in jax
 
 
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 

docs/examples_basic/second_order_problems.py

@@ -28,7 +28,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 # Quick refresher: first-order ODEs

docs/examples_basic/taylor_coefficients.py

@@ -34,7 +34,6 @@
 if not backend.has_been_selected:
     backend.select("jax")  # ivp examples in jax
 
-jax.config.update("jax_platform_name", "cpu")
 # -
 
 # We start by defining an ODE.

probdiffeq/backend/numpy.py

@@ -52,8 +52,8 @@ def flip(arr, /, axis=None):
     return jnp.flip(arr, axis=axis)
 
 
-def asarray(x, /):
-    return jnp.asarray(x)
+def asarray(x, /, dtype=None):
+    return jnp.asarray(x, dtype=dtype)
 
 
 def squeeze(arr, /):
@@ -140,10 +140,6 @@ def load(path, /):
     return jnp.load(path, allow_pickle=True)
 
 
-def allclose(a, b, *, atol=1e-8, rtol=1e-5):
-    return jnp.allclose(a, b, atol=atol, rtol=rtol)
-
-
 def stack(list_of_arrays, /):
     return jnp.stack(list_of_arrays)
 

probdiffeq/backend/testing.py

@@ -9,6 +9,7 @@
 This is not good for extendability of the test suite.
 """
 
+import jax
 import jax.numpy as jnp
 import jax.tree_util
 import pytest
@@ -31,22 +32,46 @@ def fixture(name=None, scope="function"):
     return pytest_cases.fixture(name=name, scope=scope)
 
 
-def tree_all_allclose(tree1, tree2, **kwargs):
-    trees_is_allclose = tree_allclose(tree1, tree2, **kwargs)
+def allclose(tree1, tree2, /, *, atol: float | None = None, rtol: float | None = None):
+    """Check whether two pytrees are 'close' to each other.
+
+    In contrast to jax.numpy.allclose, this version:
+    - Works with pytrees (by comparing the structure and leaves)
+    - Uses different tolerances for single precision than for double precision.
+      (It adjusts atol and rtol to the floating-point precision of the leaves.)
+    """
+    trees_is_allclose = _tree_allclose(tree1, tree2, atol=atol, rtol=rtol)
     return jax.tree_util.tree_all(trees_is_allclose)
 
 
-def tree_allclose(tree1, tree2, **kwargs):
-    def allclose_partial(*args):
-        return jnp.allclose(*args, **kwargs)
+def _tree_allclose(tree1, tree2, /, *, atol, rtol):
+    def allclose_partial(t1, t2, /):
+        return _allclose(t1, t2, atol=atol, rtol=rtol)
 
     return jax.tree_util.tree_map(allclose_partial, tree1, tree2)
 
 
-def marginals_allclose(m1, m2, /, *, ssm):
+def marginals_allclose(
+    m1, m2, /, *, ssm, atol: float | None = None, rtol: float | None = None
+):
     m1, c1 = ssm.stats.to_multivariate_normal(m1)
     m2, c2 = ssm.stats.to_multivariate_normal(m2)
 
-    means_allclose = jnp.allclose(m1, m2)
-    covs_allclose = jnp.allclose(c1, c2)
+    means_allclose = _allclose(m1, m2, atol=atol, rtol=rtol)
+    covs_allclose = _allclose(c1, c2, atol=atol, rtol=rtol)
     return jnp.logical_and(means_allclose, covs_allclose)
+
+
+def _allclose(a, b, /, *, atol: float | None, rtol: float | None):
+    # promote to float-type to enable finfo.eps
+    a = jnp.asarray(1.0 * a)
+    b = jnp.asarray(1.0 * b)
+
+    # numpy.allclose uses defaults atol=1e-8 and rtol=1e-5;
+    # we mirror this as atol=sqrt(tol) and rtol slightly larger.
+    tol = jnp.sqrt(jnp.finfo(b.dtype).eps)
+    if atol is None:
+        atol = tol
+    if rtol is None:
+        rtol = 10 * tol
+    return jnp.allclose(a, b, atol=atol, rtol=rtol)

probdiffeq/impl/_stats.py

@@ -63,16 +63,16 @@ def mahalanobis_norm_relative(self, u, /, rv):
         return np.reshape(np.abs(mahalanobis) / np.sqrt(rv.mean.size), ())
 
     def logpdf(self, u, /, rv):
-        # The cholesky factor is triangular, so we compute a cheap slogdet.
-        diagonal = linalg.diagonal_along_axis(rv.cholesky, axis1=-1, axis2=-2)
+        cholesky = linalg.qr_r(rv.cholesky.T).T
+        diagonal = linalg.diagonal_along_axis(cholesky, axis1=-1, axis2=-2)
         slogdet = np.sum(np.log(np.abs(diagonal)))
 
         dx = u - rv.mean
-        residual_white = linalg.solve_triangular(rv.cholesky.T, dx, trans="T")
-        x1 = linalg.vector_dot(residual_white, residual_white)
-        x2 = 2.0 * slogdet
-        x3 = u.size * np.log(np.pi() * 2)
-        return -0.5 * (x1 + x2 + x3)
+        residual_white = linalg.solve_triangular(cholesky, dx, lower=True, trans=0)
+        sqrnorm = linalg.vector_dot(residual_white, residual_white)
+
+        const = np.log(np.pi() * 2)
+        return -1 / 2 * sqrnorm - u.size / 2 * const - slogdet
 
     def mean(self, rv):
         return rv.mean
@@ -128,12 +128,14 @@ def logpdf(self, u, /, rv):
             u = u[None, :]
 
         def logpdf_scalar(x, r):
+            cholesky = linalg.qr_r(r.cholesky.T).T
+
             dx = x - r.mean
-            w = linalg.solve_triangular(r.cholesky.T, dx, trans="T")
+            w = linalg.solve_triangular(cholesky.T, dx, trans="T")
 
             maha_term = linalg.vector_dot(w, w)
 
-            diagonal = linalg.diagonal_along_axis(r.cholesky, axis1=-1, axis2=-2)
+            diagonal = linalg.diagonal_along_axis(cholesky, axis1=-1, axis2=-2)
             slogdet = np.sum(np.log(np.abs(diagonal)))
             logdet_term = 2.0 * slogdet
             return -0.5 * (logdet_term + maha_term + x.size * np.log(np.pi() * 2))
@@ -195,12 +197,14 @@ def mahalanobis_norm_relative(self, u, /, rv):
 
     def logpdf(self, u, /, rv):
         def logpdf_scalar(x, r):
+            cholesky = linalg.qr_r(r.cholesky.T).T
+
             dx = x - r.mean
-            w = linalg.solve_triangular(r.cholesky.T, dx, trans="T")
+            w = linalg.solve_triangular(cholesky.T, dx, trans="T")
 
             maha_term = linalg.vector_dot(w, w)
 
-            diagonal = linalg.diagonal_along_axis(r.cholesky, axis1=-1, axis2=-2)
+            diagonal = linalg.diagonal_along_axis(cholesky, axis1=-1, axis2=-2)
             slogdet = np.sum(np.log(np.abs(diagonal)))
             logdet_term = 2.0 * slogdet
             return -0.5 * (logdet_term + maha_term + x.size * np.log(np.pi() * 2))