Enabling PSR, analytical diff with pytorch and moving interfaces to dedicated module

pyproject.toml

@@ -14,10 +14,10 @@ packages = [{ include = "qiboml", from = "src" }]
 python = ">=3.9,<3.13"
 numpy = "^1.26.4"
 keras = { version = "^3.0.0", optional = true }
-tensorflow = { version = "^2.16.1", markers = "sys_platform == 'linux' or sys_platform == 'darwin'", optional = true }
+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 = { version = "^2.3.1", optional = true }
+torch = { version = "^2.3.1"}
 qibo = {git="https://github.com/qiboteam/qibo"}
 jax = "^0.4.25"
 jaxlib = "^0.4.25"
@@ -33,11 +33,10 @@ pdbpp = "^0.10.3"
 optional = true
 
 [tool.poetry.group.tests.dependencies]
-torch = "^2.3.1"
-tensorflow = { version = "^2.16.1", markers = "sys_platform == 'linux'" }
 pytest = "^7.2.1"
 pylint = "3.1.0"
 pytest-cov = "4.0.0"
+qibojit = "^0.1.7"
 
 [tool.poetry.group.benchmark.dependencies]
 pytest-benchmark = { version = "^4.0.0", extras = ["histogram"] }

src/qiboml/__init__.py

@@ -12,7 +12,7 @@
 try:
     from tensorflow import Tensor as tf_tensor
 
-    from qiboml.models import keras
+    from qiboml.interfaces import keras
 
     ndarray = Union[ndarray, tf_tensor]
 except ImportError:  # pragma: no cover
@@ -21,7 +21,7 @@
 try:
     from torch import Tensor as pt_tensor
 
-    from qiboml.models import pytorch
+    from qiboml.interfaces import pytorch
 
     ndarray = Union[ndarray, pt_tensor]
 except ImportError:  # pragma: no cover

src/qiboml/models/pytorch.py → src/qiboml/interfaces/pytorch.py

@@ -2,21 +2,13 @@
 
 from dataclasses import dataclass
 
-import numpy as np
 import torch
 from qibo import Circuit
-from qibo.backends import Backend, _check_backend
-from qibo.config import raise_error
+from qibo.backends import Backend
 
 from qiboml.models.decoding import QuantumDecoding
 from qiboml.models.encoding import QuantumEncoding
-from qiboml.operations import differentiation as Diff
-
-BACKEND_2_DIFFERENTIATION = {
-    "pytorch": None,
-    "qibolab": "PSR",
-    "jax": "Jax",
-}
+from qiboml.operations.differentiation import DifferentiationRule
 
 
 @dataclass(eq=False)
@@ -25,30 +17,23 @@ class QuantumModel(torch.nn.Module):
     encoding: QuantumEncoding
     circuit: Circuit
     decoding: QuantumDecoding
-    differentiation: str = "auto"
+    differentiation_rule: DifferentiationRule = None
 
     def __post_init__(
         self,
     ):
         super().__init__()
 
+        circuit = self.encoding.circuit
         params = [p for param in self.circuit.get_parameters() for p in param]
-        params = torch.as_tensor(self.backend.to_numpy(params)).ravel()
+        params = torch.as_tensor(self.backend.to_numpy(x=params)).ravel()
         params.requires_grad = True
         self.circuit_parameters = torch.nn.Parameter(params)
 
-        if self.differentiation == "auto":
-            self.differentiation = BACKEND_2_DIFFERENTIATION.get(
-                self.backend.name, "PSR"
-            )
-
-        if self.differentiation is not None:
-            self.differentiation = getattr(Diff, self.differentiation)()
-
     def forward(self, x: torch.Tensor):
         if (
             self.backend.name != "pytorch"
-            or self.differentiation is not None
+            or self.differentiation_rule is not None
             or not self.decoding.analytic
         ):
             x = QuantumModelAutoGrad.apply(
@@ -57,7 +42,7 @@ def forward(self, x: torch.Tensor):
                 self.circuit,
                 self.decoding,
                 self.backend,
-                self.differentiation,
+                self.differentiation_rule,
                 *list(self.parameters())[0],
             )
         else:
@@ -93,15 +78,15 @@ def forward(
         circuit: Circuit,
         decoding: QuantumDecoding,
         backend,
-        differentiation,
+        differentiation_rule,
         *parameters: list[torch.nn.Parameter],
     ):
         ctx.save_for_backward(x, *parameters)
         ctx.encoding = encoding
         ctx.circuit = circuit
         ctx.decoding = decoding
         ctx.backend = backend
-        ctx.differentiation = differentiation
+        ctx.differentiation_rule = differentiation_rule
         x_clone = x.clone().detach().cpu().numpy()
         x_clone = backend.cast(x_clone, dtype=x_clone.dtype)
         params = [
@@ -127,7 +112,7 @@ def backward(ctx, grad_output: torch.Tensor):
         ]
         grad_input, *gradients = (
             torch.as_tensor(ctx.backend.to_numpy(grad).tolist())
-            for grad in ctx.differentiation.evaluate(
+            for grad in ctx.differentiation_rule.evaluate(
                 x_clone, ctx.encoding, ctx.circuit, ctx.decoding, ctx.backend, *params
             )
         )

src/qiboml/models/ansatze.py

@@ -4,15 +4,20 @@
 from qibo import Circuit, gates
 
 
-def ReuploadingCircuit(nqubits: int, qubits: list[int] = None) -> Circuit:
+def ReuploadingCircuit(
+    nqubits: int, qubits: list[int] = None, nlayers: int = 1
+) -> Circuit:
     if qubits is None:
         qubits = list(range(nqubits))
 
     circuit = Circuit(nqubits)
-    for q in qubits:
-        circuit.add(gates.RY(q, theta=random.random() * np.pi, trainable=True))
-        circuit.add(gates.RZ(q, theta=random.random() * np.pi, trainable=True))
-    for i, q in enumerate(qubits[:-2]):
-        circuit.add(gates.CNOT(q0=q, q1=qubits[i + 1]))
-    circuit.add(gates.CNOT(q0=qubits[-1], q1=qubits[0]))
+
+    for _ in range(nlayers):
+        for q in qubits:
+            circuit.add(gates.RY(q, theta=random.random() * np.pi, trainable=True))
+            circuit.add(gates.RZ(q, theta=random.random() * np.pi, trainable=True))
+        for i, q in enumerate(qubits[:-2]):
+            circuit.add(gates.CNOT(q0=q, q1=qubits[i + 1]))
+        circuit.add(gates.CNOT(q0=qubits[-1], q1=qubits[0]))
+
     return circuit

src/qiboml/models/decoding.py

@@ -14,7 +14,7 @@ class QuantumDecoding:
 
     nqubits: int
     qubits: list[int] = None
-    nshots: int = 1000
+    nshots: int = None
     analytic: bool = True
     backend: Backend = None
     _circuit: Circuit = None
@@ -58,7 +58,6 @@ def output_shape(self):
 class Expectation(QuantumDecoding):
 
     observable: Union[ndarray, Hamiltonian] = None
-    analytic: bool = False
 
     def __post_init__(self):
         if self.observable is None:
@@ -69,7 +68,7 @@ def __post_init__(self):
         super().__post_init__()
 
     def __call__(self, x: Circuit) -> ndarray:
-        if self.analytic:
+        if self.nshots is None:
             return self.observable.expectation(
                 super().__call__(x).state(),
             ).reshape(1, 1)