Compute Backends and Automatic Differentiation

.github/CODEOWNERS

@@ -1,11 +1,18 @@
 * @probabilistic-numerics/probnum-global-codeowners
 
+# Compute Backends
+/src/probnum/backend @marvinpfoertner @JonathanWenger
+/tests/probnum/backend @marvinpfoertner @JonathanWenger
+
+# Compatibility Functions
+/src/probnum/compat @marvinpfoertner @JonathanWenger
+
 # Differential Equations
 /src/probnum/diffeq/ @pnkraemer @schmidtjonathan
 /src/probnum/problems/zoo/diffeq/ @pnkraemer @schmidtjonathan
 
 /tests/test_diffeq/ @pnkraemer @schmidtjonathan
-/tests/test_problems/test_zoo/test_diffeq/ @pnkraemer @schmidtjonathan
+/tests/probnum/problems/zoo/diffeq/ @pnkraemer @schmidtjonathan
 
 /benchmarks/ivpsolvers.py @pnkraemer @schmidtjonathan
 
@@ -14,22 +21,22 @@
 /src/probnum/problems/zoo/filtsmooth/ @pnkraemer @schmidtjonathan
 
 /tests/test_filtsmooth/ @pnkraemer @schmidtjonathan
-/tests/test_problems/test_zoo/test_filtsmooth/ @pnkraemer @schmidtjonathan
+/tests/problems/zoo/filtsmooth/ @pnkraemer @schmidtjonathan
 
 /benchmarks/filtsmooth.py @pnkraemer @schmidtjonathan
 
 # Linear Algebra
 /src/probnum/linalg/ @JonathanWenger @marvinpfoertner
 /src/probnum/problems/zoo/linalg/ @JonathanWenger @marvinpfoertner
 
-/tests/test_linalg/ @JonathanWenger @marvinpfoertner
-/tests/test_problems/test_zoo/test_linalg/ @JonathanWenger @marvinpfoertner
+/tests/probnum/linalg/ @JonathanWenger @marvinpfoertner
+/tests/problems/zoo/linalg/ @JonathanWenger @marvinpfoertner
 
 /benchmarks/linearsolvers.py @JonathanWenger @marvinpfoertner
 
 # Linear Operators
 /src/probnum/linops/ @marvinpfoertner @JonathanWenger
-/tests/test_linops/ @marvinpfoertner @JonathanWenger
+/tests/probnum/linops/ @marvinpfoertner @JonathanWenger
 
 /benchmarks/linops.py @marvinpfoertner @JonathanWenger
 
@@ -38,24 +45,20 @@
 /src/probnum/problems/zoo/quad/ @mmahsereci @tskarvone
 
 /tests/test_quad/ @mmahsereci @tskarvone
-/tests/test_problems/test_zoo/test_quad/ @mmahsereci @tskarvone
+/tests/problems/zoo/quad/ @mmahsereci @tskarvone
 
 # Random Processes & Kernels
 /src/probnum/randprocs/ @marvinpfoertner @JonathanWenger
-/tests/test_randprocs/ @marvinpfoertner @JonathanWenger
+/tests/probnum/randprocs/ @marvinpfoertner @JonathanWenger
 
 /benchmarks/randprocs.py @marvinpfoertner @JonathanWenger
 /benchmarks/kernels.py @marvinpfoertner @JonathanWenger
 
 /src/probnum/randprocs/markov/ @pnkraemer @schmidtjonathan
-/tests/test_randprocs/test_markov/ @pnkraemer @schmidtjonathan
+/tests/probnum/randprocs/markov/ @pnkraemer @schmidtjonathan
 
 # Random Variables
 /src/probnum/randvars/ @marvinpfoertner @JonathanWenger
-/tests/test_randvars/ @marvinpfoertner @JonathanWenger
+/tests/probnum/randvars/ @marvinpfoertner @JonathanWenger
 
 /benchmarks/random_variables.py @marvinpfoertner @JonathanWenger
-
-# Utils
-/src/probnum/utils/linalg/_cholesky_updates.py @pnkraemer
-/tests/test_utils/test_linalg/test_cholesky_updates.py @pnkraemer

.github/workflows/CI-build.yml

@@ -16,6 +16,7 @@ jobs:
       matrix:
         platform: [ubuntu-latest, macos-latest, windows-latest]
         python: ["3.8", "3.9", "3.10"]
+        backend: ["numpy", "jax", "torch"]
 
     steps:
       - uses: actions/checkout@v2
@@ -28,11 +29,11 @@ jobs:
       - name: Install Tox and any other packages
         run: pip install tox
       - name: Run Tox
-        # Run tox using the version of Python in `PATH`
-        run: tox -e py3
+        # Run tox using the version of Python in `PATH` and the corresponding compute backend
+        run: tox -e py3-${{ matrix.backend }}
       - name: Upload coverage report to Codecov
         uses: codecov/codecov-action@v2
-        if: startsWith(matrix.platform,'ubuntu') && matrix.python == '3.8'
+        if: startsWith(matrix.platform,'ubuntu') && matrix.python == '3.8' && matrix.backend == 'numpy'
 
   documentation:
     runs-on: ubuntu-latest

benchmarks/linearsolvers.py

@@ -1,7 +1,7 @@
 """Benchmarks for linear solvers."""
 import numpy as np
 
-from probnum import linops, problems, randvars
+from probnum import backend, linops, problems, randvars
 from probnum.linalg import problinsolve
 from probnum.problems.zoo.linalg import random_sparse_spd_matrix, random_spd_matrix
 
@@ -11,26 +11,36 @@
 
 
 def get_linear_system(name: str, dim: int):
-    rng = np.random.default_rng(0)
+    rng_state = backend.random.rng_state(42)
 
     if name == "dense":
         if dim > 1000:
             raise NotImplementedError()
-        A = random_spd_matrix(rng=rng, dim=dim)
+        rng_state, rng_state_A = backend.random.split(rng_state, 2)
+        A = random_spd_matrix(rng_state=rng_state_A, shape=(dim, dim))
     elif name == "sparse":
+        rng_state, rng_state_A_sparse = backend.random.split(rng_state, 2)
         A = random_sparse_spd_matrix(
-            rng=rng, dim=dim, density=np.minimum(1.0, 1000 / dim**2)
+            rng_state=rng_state_A_sparse,
+            shape=(dim, dim),
+            density=backend.minimum(1.0, 1000 / dim**2),
         )
     elif name == "linop":
         if dim > 100:
             raise NotImplementedError()
             # TODO: Larger benchmarks currently fail. Remove once PLS refactor
             # (https://github.com/probabilistic-numerics/probnum/issues/51) is resolved
-        A = linops.Scaling(factors=rng.normal(size=(dim,)))
+        rng_state, rng_state_A_linop = backend.random.split(rng_state, 2)
+        A = linops.Scaling(
+            factors=backend.random.standard_normal(
+                rng_state=rng_state_A_linop, shape=(dim,)
+            )
+        )
     else:
         raise NotImplementedError()
 
-    solution = rng.normal(size=(dim,))
+    rng_state, rng_state_solution = backend.random.split(rng_state, 2)
+    solution = backend.random.standard_normal(rng_state_solution, shape=(dim,))
     b = A @ solution
     return problems.LinearSystem(A=A, b=b, solution=solution)
 
@@ -72,14 +82,16 @@ def peakmem_solve(self, linsys, dim):
         problinsolve(A=self.linsys.A, b=self.linsys.b)
 
     def track_residual_norm(self, linsys, dim):
-        return np.linalg.norm(self.linsys.b - self.linsys.A @ self.xhat.mean)
+        return backend.linalg.vector_norm(
+            self.linsys.b - self.linsys.A @ self.xhat.mean
+        ).item()
 
     def track_error_2norm(self, linsys, dim):
-        return np.linalg.norm(self.linsys.solution - self.xhat.mean)
+        return backend.linalg.vector_norm(self.linsys.solution - self.xhat.mean).item()
 
     def track_error_Anorm(self, linsys, dim):
         diff = self.linsys.solution - self.xhat.mean
-        return np.sqrt(np.inner(diff, self.linsys.A @ diff))
+        return backend.sqrt(np.inner(diff, self.linsys.A @ diff))
 
 
 class PosteriorBelief:

benchmarks/random_variables.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 
-from probnum import linops, randvars as rvs
+from probnum import backend, linops, randvars
 
 # Module level variables
 RV_NAMES = [
@@ -39,15 +39,15 @@ def get_randvar(rv_name):
     cov_2d_symkron = linops.SymmetricKronecker(A=SPD_MATRIX_5x5)
 
     if rv_name == "univar_normal":
-        randvar = rvs.Normal(mean=mean_0d, cov=cov_0d)
+        randvar = randvars.Normal(mean=mean_0d, cov=cov_0d)
     elif rv_name == "multivar_normal":
-        randvar = rvs.Normal(mean=mean_1d, cov=cov_1d)
+        randvar = randvars.Normal(mean=mean_1d, cov=cov_1d)
     elif rv_name == "matrixvar_normal":
-        randvar = rvs.Normal(mean=mean_2d_mat, cov=cov_2d_kron)
+        randvar = randvars.Normal(mean=mean_2d_mat, cov=cov_2d_kron)
     elif rv_name == "symmatrixvar_normal":
-        randvar = rvs.Normal(mean=mean_2d_mat, cov=cov_2d_symkron)
+        randvar = randvars.Normal(mean=mean_2d_mat, cov=cov_2d_symkron)
     elif rv_name == "operatorvar_normal":
-        randvar = rvs.Normal(mean=mean_2d_linop, cov=cov_2d_symkron)
+        randvar = randvars.Normal(mean=mean_2d_linop, cov=cov_2d_symkron)
     else:
         raise ValueError("Random variable not found.")
 
@@ -87,14 +87,14 @@ class Sampling:
     params = [RV_NAMES]
 
     def setup(self, randvar):
-        self.rng = np.random.default_rng(seed=2)
+        self.rng_state = backend.random.rng_state(23529)
         self.n_samples = 1000
         self.randvar = get_randvar(rv_name=randvar)
 
     def time_sample(self, randvar):
         """Times sampling from this distribution."""
-        self.randvar.sample(rng=self.rng, size=self.n_samples)
+        self.randvar.sample(rng_state=self.rng_state, sample_shape=self.n_samples)
 
     def peakmem_sample(self, randvar):
         """Peak memory of sampling process."""
-        self.randvar.sample(rng=self.rng, size=self.n_samples)
+        self.randvar.sample(rng_state=self.rng_state, sample_shape=self.n_samples)

docs/source/api.rst

@@ -7,6 +7,10 @@ API Reference
     +-------------------------------------------------+--------------------------------------------------------------+
     | **Subpackage**                                  | **Description**                                              |
     +-------------------------------------------------+--------------------------------------------------------------+
+    | :mod:`~probnum.backend`                         | Generic computation backend.                                 |
+    +-------------------------------------------------+--------------------------------------------------------------+
+    | :mod:`~probnum.compat`                          | Compatibility functions.                                     |
+    +-------------------------------------------------+--------------------------------------------------------------+
     | :class:`config <probnum._config.Configuration>` | Global configuration options.                                |
     +-------------------------------------------------+--------------------------------------------------------------+
     | :mod:`~probnum.diffeq`                          | Probabilistic solvers for ordinary differential equations.   |
@@ -29,15 +33,15 @@ API Reference
     +-------------------------------------------------+--------------------------------------------------------------+
     | :mod:`~probnum.typing`                          | Type aliases.                                                |
     +-------------------------------------------------+--------------------------------------------------------------+
-    | :mod:`~probnum.utils`                           | Utility functions.                                           |
-    +-------------------------------------------------+--------------------------------------------------------------+
 
 
 .. toctree::
     :caption: API Reference
     :hidden:
 
     api/probnum
+    api/backend
+    api/compat
     api/config
     api/diffeq
     api/filtsmooth
@@ -49,4 +53,3 @@ API Reference
     api/randprocs
     api/randvars
     api/typing
-    api/utils

docs/source/api/backend.rst

@@ -0,0 +1,100 @@
+***************
+probnum.backend
+***************
+
+.. automodule:: probnum.backend
+
+.. currentmodule:: probnum.backend
+
+Classes
+-------
+
+.. autosummary::
+
+    ~probnum.backend.Dispatcher
+
+
+.. toctree::
+    :hidden:
+
+    backend/array_object
+
+.. toctree::
+    :hidden:
+
+    backend/data_types
+
+.. toctree::
+    :hidden:
+
+    backend/creation_functions
+
+.. toctree::
+    :hidden:
+
+    backend/elementwise_functions
+
+.. toctree::
+    :hidden:
+
+    backend/logic_functions
+
+.. toctree::
+    :hidden:
+
+    backend/manipulation_functions
+
+.. toctree::
+    :hidden:
+
+    backend/searching_functions
+
+.. toctree::
+    :hidden:
+
+    backend/sorting_functions
+
+.. toctree::
+    :hidden:
+
+    backend/statistical_functions
+
+.. toctree::
+    :hidden:
+
+    backend/jit_compilation
+
+.. toctree::
+    :hidden:
+
+    backend/vectorization
+
+.. toctree::
+    :hidden:
+
+    backend/probnum.backend.Dispatcher
+
+.. toctree::
+    :hidden:
+
+    backend/autodiff
+
+.. toctree::
+    :hidden:
+
+    backend/linalg
+
+.. toctree::
+    :hidden:
+
+    backend/random
+
+.. toctree::
+    :hidden:
+
+    backend/special
+
+.. toctree::
+    :hidden:
+
+    backend/typing

docs/source/api/backend/array_object.rst

@@ -0,0 +1,40 @@
+Array Object
+============
+
+The basic object representing a multi-dimensional array and adjacent functionality.
+
+.. currentmodule:: probnum.backend
+
+Functions
+---------
+
+.. autosummary::
+
+    ~probnum.backend.asshape
+    ~probnum.backend.isarray
+    ~probnum.backend.ndim
+    ~probnum.backend.to_numpy
+
+Classes
+-------
+
++----------------------------------+---------------------------------------------------------------------------------------------------------------------------------+
+| :class:`~probnum.backend.Array`  | Object representing a multi-dimensional array stored on a :class:`~probnum.backend.Device` and containing elements of the same  |
+|                                  | :class:`~probnum.backend.DType`.                                                                                                |
++----------------------------------+---------------------------------------------------------------------------------------------------------------------------------+
+| :class:`~probnum.backend.Scalar` | Object representing a scalar with a :class:`~probnum.backend.DType`.                                                            |
++----------------------------------+---------------------------------------------------------------------------------------------------------------------------------+
+| :class:`~probnum.backend.Device` | Device, such as a CPU or GPU, on which an :class:`~probnum.backend.Array` is located.                                           |
++----------------------------------+---------------------------------------------------------------------------------------------------------------------------------+
+
+
+.. toctree::
+    :hidden:
+
+    array_object/probnum.backend.asshape
+    array_object/probnum.backend.isarray
+    array_object/probnum.backend.ndim
+    array_object/probnum.backend.to_numpy
+    array_object/probnum.backend.Array
+    array_object/probnum.backend.Device
+    array_object/probnum.backend.Scalar