Add einsum

Author: marvinpfoertner

src/probnum/backend/__init__.py

@@ -1,52 +1,5 @@
 from ._select import Backend, select_backend as _select_backend
 
-# pylint: disable=undefined-all-variable
-__all__ = [
-    "ndarray",
-    # DTypes
-    "dtype",
-    "asdtype",
-    "bool",
-    "int32",
-    "int64",
-    "single",
-    "double",
-    "csingle",
-    "cdouble",
-    "cast",
-    "promote_types",
-    "is_floating",
-    "finfo",
-    # Shape Arithmetic
-    "reshape",
-    "atleast_1d",
-    "atleast_2d",
-    "broadcast_arrays",
-    "broadcast_shapes",
-    "ndim",
-    "swapaxes",
-    # Constructors
-    "array",
-    "asarray",
-    "diag",
-    "eye",
-    "ones",
-    "ones_like",
-    "zeros",
-    "zeros_like",
-    "linspace",
-    # Constants
-    "pi",
-    "inf",
-    # Operations
-    "sin",
-    "exp",
-    "log",
-    "sqrt",
-    "sum",
-    "maximum",
-]
-
 BACKEND = _select_backend()
 
 # isort: off
@@ -56,10 +9,17 @@
 from ._core import *
 
 from . import (
+    _core,
     autodiff,
     linalg,
     random,
     special,
 )
 
 # isort: on
+
+__all__ = [
+    "Backend",
+    "BACKEND",
+    "Dispatcher",
+] + _core.__all__

src/probnum/backend/_core/__init__.py

@@ -64,6 +64,9 @@
 # Element-wise Binary Operations
 maximum = _core.maximum
 
+# Contractions
+einsum = _core.einsum
+
 # Reductions
 all = _core.all
 sum = _core.sum
@@ -91,3 +94,64 @@ def as_scalar(x: ScalarArgType, dtype: DTypeArgType = None) -> ArrayType:
         raise ValueError("The given input is not a scalar.")
 
     return asarray(x, dtype=dtype)[()]
+
+
+__all__ = [
+    "ndarray",
+    # DTypes
+    "dtype",
+    "asdtype",
+    "bool",
+    "int32",
+    "int64",
+    "single",
+    "double",
+    "csingle",
+    "cdouble",
+    "cast",
+    "promote_types",
+    "is_floating",
+    "finfo",
+    # Shape Arithmetic
+    "reshape",
+    "atleast_1d",
+    "atleast_2d",
+    "broadcast_arrays",
+    "broadcast_shapes",
+    "ndim",
+    "swapaxes",
+    # Constructors
+    "array",
+    "asarray",
+    "as_scalar",
+    "diag",
+    "eye",
+    "full",
+    "full_like",
+    "ones",
+    "ones_like",
+    "zeros",
+    "zeros_like",
+    "linspace",
+    # Constants
+    "inf",
+    "pi",
+    # Element-wise Unary Operations
+    "exp",
+    "isfinite",
+    "log",
+    "sin",
+    "sqrt",
+    # Element-wise Binary Operations
+    "maximum",
+    # Contractions
+    "einsum",
+    # Reductions
+    "all",
+    "sum",
+    # Misc
+    "to_numpy",
+    # Just-in-Time Compilation
+    "jit",
+    "jit_method",
+]

src/probnum/backend/_core/_jax.py

@@ -17,6 +17,7 @@
     double,
     dtype,
     dtype as asdtype,
+    einsum,
     exp,
     eye,
     finfo,

src/probnum/backend/_core/_numpy.py

@@ -16,6 +16,7 @@
     double,
     dtype,
     dtype as asdtype,
+    einsum,
     exp,
     eye,
     finfo,

src/probnum/backend/_core/_torch.py

@@ -15,6 +15,7 @@
     diag,
     double,
     dtype,
+    einsum,
     exp,
     eye,
     finfo,