Merge pull request #172 from qiboteam/list_backends

MetaBackend

src/qibojit/__init__.py

@@ -1,3 +1,5 @@
 import importlib.metadata as im
 
+from qibojit.backends import MetaBackend
+
 __version__ = im.version(__package__)

src/qibojit/backends/__init__.py

@@ -1,2 +1,46 @@
+from typing import Union
+
 from qibojit.backends.cpu import NumbaBackend
 from qibojit.backends.gpu import CupyBackend, CuQuantumBackend
+
+QibojitBackend = Union[NumbaBackend, CupyBackend, CuQuantumBackend]
+
+PLATFORMS = ("numba", "cupy", "cuquantum")
+
+
+class MetaBackend:
+    """Meta-backend class which takes care of loading the qibojit backends."""
+
+    @staticmethod
+    def load(platform: str = None) -> QibojitBackend:
+        """Loads the backend.
+
+        Args:
+            platform (str): Name of the backend to load: either `numba`, `cupy` or `cuquantum`.
+        Returns:
+            qibo.backends.abstract.Backend: The loaded backend.
+        """
+
+        if platform == "numba":
+            return NumbaBackend()
+        elif platform == "cupy":
+            return CupyBackend()
+        elif platform == "cuquantum":
+            return CuQuantumBackend()
+        else:  # pragma: no cover
+            try:
+                return CupyBackend()
+            except (ModuleNotFoundError, ImportError):
+                return NumbaBackend()
+
+    def list_available(self) -> dict:
+        """Lists all the available qibojit backends."""
+        available_backends = {}
+        for platform in PLATFORMS:
+            try:
+                MetaBackend.load(platform=platform)
+                available = True
+            except:
+                available = False
+            available_backends[platform] = available
+        return available_backends

src/qibojit/tests/test_backends.py

@@ -1,6 +1,8 @@
 import numpy as np
 import pytest
 
+from qibojit.backends import MetaBackend
+
 
 def test_device_setter(backend):
     if backend.platform == "numba":
@@ -114,3 +116,8 @@ def test_backend_eigh_sparse(backend, sparse_type, k):
     eigvals1 = backend.to_numpy(eigvals1)
     eigvals2 = backend.to_numpy(eigvals2)
     backend.assert_allclose(sorted(eigvals1), sorted(eigvals2))
+
+
+def test_metabackend_list_available():
+    available_backends = dict(zip(("numba", "cupy", "cuquantum"), (True, False, False)))
+    assert MetaBackend().list_available() == available_backends