Merge pull request #37 from qiboteam/pytorch_autodiff

Pytorch interface autodifferentiation

Author: BrunoLiegiBastonLiegi

poetry.lock

@@ -18,7 +18,7 @@ tensorflow = { version = "^2.16.1", markers = "sys_platform == 'linux' or sys_pl
 # 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 }
-qibo = {git="https://github.com/qiboteam/qibo", branch="qiboml_models_updates"}
+qibo = {git="https://github.com/qiboteam/qibo"}
 jax = "^0.4.25"
 jaxlib = "^0.4.25"
 

pyproject.toml

@@ -20,6 +20,7 @@ def __init__(self):
 
         self.np = jnp
         self.tensor_types = (jnp.ndarray, numpy.ndarray)
+        self.matrices.np = jnp
 
     def set_precision(self, precision):
         if precision != self.precision:
@@ -43,6 +44,13 @@ def cast(self, x, dtype=None, copy=False):
             return x.astype(dtype)
         return self.np.array(x, dtype=dtype, copy=copy)
 
+    def to_numpy(self, x):
+
+        if isinstance(x, list) or isinstance(x, tuple):
+            return self.numpy.asarray([self.to_numpy(i) for i in x])
+
+        return self.numpy.asarray(x)
+
     # TODO: using numpy's rng for now. Shall we use Jax's?
     def set_seed(self, seed):
         self.numpy.random.seed(seed)

src/qiboml/backends/jax.py

@@ -1,22 +1,18 @@
-from dataclasses import dataclass
+import random
 
 import numpy as np
 from qibo import Circuit, gates
 
-from qiboml.models.abstract import QuantumCircuitLayer
 
+def ReuploadingCircuit(nqubits: int, qubits: list[int] = None) -> Circuit:
+    if qubits is None:
+        qubits = list(range(nqubits))
 
-@dataclass
-class ReuploadingLayer(QuantumCircuitLayer):
-
-    def __post_init__(self):
-        super().__post_init__()
-        for q in self.qubits:
-            self.circuit.add(gates.RY(q, theta=0.0))
-            self.circuit.add(gates.RZ(q, theta=0.0))
-        for i, q in enumerate(self.qubits[:-2]):
-            self.circuit.add(gates.CNOT(q0=q, q1=self.qubits[i + 1]))
-        self.circuit.add(gates.CNOT(q0=self.qubits[-1], q1=self.qubits[0]))
-
-    def forward(self, x: Circuit) -> Circuit:
-        return x + self.circuit
+    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]))
+    return circuit

src/qiboml/models/abstract.py

@@ -0,0 +1,117 @@
+from dataclasses import dataclass
+from typing import Union
+
+from qibo import Circuit, gates
+from qibo.backends import Backend, _check_backend
+from qibo.config import raise_error
+from qibo.hamiltonians import Hamiltonian
+
+from qiboml import ndarray
+
+
+@dataclass
+class QuantumDecoding:
+
+    nqubits: int
+    qubits: list[int] = None
+    nshots: int = 1000
+    analytic: bool = True
+    backend: Backend = None
+    _circuit: Circuit = None
+
+    def __post_init__(self):
+        if self.qubits is None:
+            self.qubits = list(range(self.nqubits))
+        self._circuit = Circuit(self.nqubits)
+        self.backend = _check_backend(self.backend)
+        self._circuit.add(gates.M(*self.qubits))
+
+    def __call__(self, x: Circuit) -> "CircuitResult":
+        return self.backend.execute_circuit(x + self._circuit, nshots=self.nshots)
+
+    @property
+    def circuit(
+        self,
+    ):
+        return self._circuit
+
+    def set_backend(self, backend):
+        self.backend = backend
+
+    @property
+    def output_shape(self):
+        raise_error(NotImplementedError)
+
+
+@dataclass
+class Probabilities(QuantumDecoding):
+
+    def __call__(self, x: Circuit) -> ndarray:
+        return super().__call__(x).probabilities()
+
+    @property
+    def output_shape(self):
+        return (1, 2**self.nqubits)
+
+
+@dataclass
+class Expectation(QuantumDecoding):
+
+    observable: Union[ndarray, Hamiltonian] = None
+    analytic: bool = False
+
+    def __post_init__(self):
+        if self.observable is None:
+            raise_error(
+                RuntimeError,
+                "Please provide an observable for expectation value calculation.",
+            )
+        super().__post_init__()
+
+    def __call__(self, x: Circuit) -> ndarray:
+        if self.analytic:
+            return self.observable.expectation(
+                super().__call__(x).state(),
+            ).reshape(1, 1)
+        else:
+            return self.observable.expectation_from_samples(
+                super().__call__(x).frequencies(),
+                qubit_map=self.qubits,
+            ).reshape(1, 1)
+
+    @property
+    def output_shape(self):
+        return (1, 1)
+
+    def set_backend(self, backend):
+        super().set_backend(backend)
+        self.observable.backend = backend
+
+
+@dataclass
+class State(QuantumDecoding):
+
+    def __call__(self, x: Circuit) -> ndarray:
+        state = super().__call__(x).state()
+        return self.backend.np.vstack(
+            (self.backend.np.real(state), self.backend.np.imag(state))
+        ).reshape(self.output_shape)
+
+    @property
+    def output_shape(self):
+        return (2, 1, 2**self.nqubits)
+
+
+@dataclass
+class Samples(QuantumDecoding):
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.analytic = False
+
+    def forward(self, x: Circuit) -> ndarray:
+        return self.backend.cast(super().__call__(x).samples(), self.backend.precision)
+
+    @property
+    def output_shape(self):
+        return (self.nshots, len(self.qubits))

src/qiboml/models/ansatze.py

@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import numpy as np
+from qibo import Circuit, gates
+from qibo.config import raise_error
+
+from qiboml import ndarray
+
+
+@dataclass
+class QuantumEncoding(ABC):
+
+    nqubits: int
+    qubits: list[int] = None
+    _circuit: Circuit = None
+
+    def __post_init__(
+        self,
+    ):
+        if self.qubits is None:
+            self.qubits = list(range(self.nqubits))
+        self._circuit = Circuit(self.nqubits)
+
+    @abstractmethod
+    def __call__(self, x: ndarray) -> Circuit:
+        pass
+
+    @property
+    def circuit(
+        self,
+    ):
+        return self._circuit
+
+
+@dataclass
+class PhaseEncoding(QuantumEncoding):
+
+    def __post_init__(
+        self,
+    ):
+        super().__post_init__()
+        for q in self.qubits:
+            self._circuit.add(gates.RY(q, theta=0.0, trainable=False))
+
+    def _set_phases(self, x: ndarray):
+        for gate, phase in zip(self._circuit.parametrized_gates, x.ravel()):
+            gate.parameters = phase
+
+    def __call__(self, x: ndarray) -> Circuit:
+        self._set_phases(x)
+        return self._circuit
+
+
+@dataclass
+class BinaryEncoding(QuantumEncoding):
+
+    def __call__(self, x: ndarray) -> Circuit:
+        if x.shape[-1] != len(self.qubits):
+            raise_error(
+                RuntimeError,
+                f"Invalid input dimension {x.shape[-1]}, but the allocated qubits are {self.qubits}.",
+            )
+        circuit = self.circuit.copy()
+        ones = np.flatnonzero(x.ravel() == 1)
+        for bit in ones:
+            circuit.add(gates.X(self.qubits[bit]))
+        return circuit