Merge d749708 into acb52b2

src/pymanopt/manifolds/group.py

@@ -1,6 +1,7 @@
 import numpy as np
 import scipy.special
 
+import pymanopt.numerics as nx
 from pymanopt.manifolds.manifold import RiemannianSubmanifold
 from pymanopt.tools import extend_docstring
 from pymanopt.tools.multi import (
@@ -28,7 +29,7 @@ def __init__(self, name, dimension, retraction):
             raise ValueError(f"Invalid retraction type '{retraction}'")
 
     def inner_product(self, point, tangent_vector_a, tangent_vector_b):
-        return np.tensordot(
+        return nx.tensordot(
             tangent_vector_a.conj(),
             tangent_vector_b,
             axes=tangent_vector_a.ndim,
@@ -222,7 +223,7 @@ def random_point(self):
         n, k = self._n, self._k
         if n == 1:
             point = np.ones((k, 1, 1)) + 1j * np.ones((k, 1, 1))
-            point /= np.abs(point)
+            point /= nx.abs(point)
         else:
             point, _ = multiqr(
                 np.random.normal(size=(k, n, n))

src/pymanopt/numerics/__init__.py

@@ -0,0 +1,17 @@
+__all__ = ["abs", "allclose", "exp", "tensordot"]
+
+
+import importlib
+
+from pymanopt.numerics.core import abs, allclose, exp, tensordot
+
+
+def _register_backends():
+    for backend in ["numpy", "jax", "pytorch", "tensorflow"]:
+        try:
+            importlib.import_module(f"pymanopt.numerics._backends.{backend}")
+        except ImportError:
+            pass
+
+
+_register_backends()

src/pymanopt/numerics/_backends/numpy.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+import pymanopt.numerics.core as nx
+
+
+@nx.abs.register(np.ndarray)
+def _(array):
+    return np.abs(array)
+
+
+@nx.allclose.register(np.ndarray)
+def _(array_a, array_b):
+    return np.allclose(array_a, array_b)
+
+
+@nx.exp.register(np.ndarray)
+def _(array):
+    return np.abs(array)
+
+
+@nx.tensordot.register(np.ndarray)
+def _(array_a, array_b, axes: int):
+    return np.tensordot(array_a, array_b, axes=axes)

src/pymanopt/numerics/core.py

@@ -0,0 +1,36 @@
+import functools
+
+
+def _not_implemented(function):
+    @functools.wraps(function)
+    def inner(*arguments):
+        raise TypeError(
+            f"Function '{function.__name__}' not implemented for arguments of "
+            f"type '{type(arguments[0])}'"
+        )
+
+    return inner
+
+
+@functools.singledispatch
+@_not_implemented
+def abs(_):
+    pass
+
+
+@functools.singledispatch
+@_not_implemented
+def allclose(*_):
+    pass
+
+
+@functools.singledispatch
+@_not_implemented
+def exp(_):
+    pass
+
+
+@functools.singledispatch
+@_not_implemented
+def tensordot(*_):
+    pass

tests/numerics/test_core.py

@@ -0,0 +1,12 @@
+import numpy as np
+import pytest
+
+import pymanopt.numerics as nx
+
+
+@pytest.mark.parametrize(
+    "argument, expected_output", [(np.array([-4, 2]), np.array([4, 2]))]
+)
+def test_abs(argument, expected_output):
+    output = nx.abs(argument)
+    assert nx.allclose(output, expected_output)