Symbolical differentiation using TensorFlow

.github/workflows/.codecov.yml

@@ -0,0 +1,29 @@
+# Reference docs at:
+# https://docs.codecov.io/docs/codecovyml-reference
+
+codecov:
+  require_ci_to_pass: yes
+coverage:
+  precision: 2
+  round: down
+  range: "70...100"
+  status:
+    project:
+      default:
+        threshold: 1%
+
+parsers:
+  gcov:
+    branch_detection:
+      conditional: yes
+      loop: yes
+      method: no
+      macro: no
+
+comment:
+  layout: "reach,diff,flags,tree"
+  behavior: default
+  require_changes: no
+
+github_checks:
+  annotations: false

pyproject.toml

@@ -13,13 +13,12 @@ packages = [{ include = "qiboml", from = "src" }]
 [tool.poetry.dependencies]
 python = ">=3.9,<3.12"
 numpy = "^1.26.4"
-numba = "^0.59.0"
 tensorflow = { version = "^2.16.1", markers = "sys_platform == 'linux' or sys_platform == 'darwin'" }
 # TODO: the marker is a temporary solution due to the lack of the tensorflow-io 0.32.0's wheels for Windows, this package is one of
 # the tensorflow requirements
 torch = "^2.2.0"
 jax = "^0.4.25"
-qibo = "^0.2.6"
+qibo = "^0.2.8"
 
 [tool.poetry.group.dev]
 optional = true
@@ -35,6 +34,7 @@ pytest-cov = "^3.0.0"
 pytest-env = "^0.8.1"
 pytest-benchmark = { version = "^4.0.0", extras = ["histogram"] }
 
+
 [tool.poe.tasks]
 bench = "pytest benchmarks/"
 test = "pytest"

src/qiboml/backends/__init__.py

@@ -1 +1,3 @@
-from qiboml.backends.tensorflow import TensorflowBackend
+from qiboml.backends.jax import JaxBackend
+from qiboml.backends.pytorch import PyTorchBackend
+from qiboml.backends.tensorflow import TensorflowBackend

src/qiboml/backends/jax.py

@@ -1,25 +1,50 @@
 from qibo import __version__
 from qibo.backends import einsum_utils
+from qibo.backends.npmatrices import NumpyMatrices
 from qibo.backends.numpy import NumpyBackend
 from qibo.config import raise_error
 
 
+class JaxMatrices(NumpyMatrices):
+    def __init__(self, dtype):
+        super().__init__(dtype)
+        import jax  # pylint: disable=import-error
+        import jax.numpy as jnp  # pylint: disable=import-error
+
+        self.jax = jax
+        self.np = jnp
+
+    def _cast(self, x, dtype):
+        return self.np.array(x, dtype=dtype)
+
+    def Unitary(self, u):
+        return self._cast(u, dtype=self.dtype)
+
+
 class JaxBackend(NumpyBackend):
     def __init__(self):
         super().__init__()
         self.name = "jax"
 
         import jax
         import jax.numpy as jnp  # pylint: disable=import-error
-        import numpy
+        import numpy as np
 
         jax.config.update("jax_enable_x64", True)
 
         self.jax = jax
-        self.numpy = numpy
+        self.numpy = np
 
         self.np = jnp
-        self.tensor_types = (jnp.ndarray, numpy.ndarray)
+
+        self.versions = {
+            "qibo": __version__,
+            "numpy": np.__version__,
+            "tensorflow": jax.__version__,
+        }
+
+        self.matrices = JaxMatrices(self.dtype)
+        self.tensor_types = (jnp.ndarray, np.ndarray)
 
     def set_precision(self, precision):
         if precision != self.precision:
@@ -72,6 +97,18 @@ def plus_density_matrix(self, nqubits):
         state /= 2**nqubits
         return state
 
+    def matrix(self, gate):
+        npmatrix = super().matrix(gate)
+        return self.np.array(npmatrix, dtype=self.dtype)
+
+    def matrix_parametrized(self, gate):
+        npmatrix = super().matrix_parametrized(gate)
+        return self.np.array(npmatrix, dtype=self.dtype)
+
+    def matrix_fused(self, gate):
+        npmatrix = super().matrix_fused(gate)
+        return self.np.array(npmatrix, dtype=self.dtype)
+
     def update_frequencies(self, frequencies, probabilities, nsamples):
         samples = self.sample_shots(probabilities, nsamples)
         res, counts = self.np.unique(samples, return_counts=True)

src/qiboml/backends/tensorflow.py

@@ -31,9 +31,13 @@ def __init__(self):
         super().__init__()
         self.name = "tensorflow"
         os.environ["TF_CPP_MIN_LOG_LEVEL"] = str(TF_LOG_LEVEL)
+
         import tensorflow as tf  # pylint: disable=import-error
         import tensorflow.experimental.numpy as tnp  # pylint: disable=import-error
 
+        if TF_LOG_LEVEL >= 2:
+            tf.get_logger().setLevel("ERROR")
+
         tnp.experimental_enable_numpy_behavior()
         self.tf = tf
         self.np = tnp