fix: various fixes to keras interface + testing both keras and torch …

…interfaces with the frontend fixture

pyproject.toml

@@ -33,12 +33,12 @@ optional = true
 [tool.poetry.group.tests.dependencies]
 torch = "^2.3.1"
 tensorflow = { version = "^2.16.1", markers = "sys_platform == 'linux'" }
-
-[tool.poetry.group.benchmark.dependencies]
 pytest = "^7.1.2"
 pylint = "^2.17"
 pytest-cov = "^3.0.0"
 pytest-env = "^0.8.1"
+
+[tool.poetry.group.benchmark.dependencies]
 pytest-benchmark = { version = "^4.0.0", extras = ["histogram"] }
 
 [tool.poe.tasks]

src/qiboml/models/encoding_decoding.py

@@ -76,15 +76,26 @@ def backward(self):
 
 class ProbabilitiesLayer(QuantumDecodingLayer):
 
+    def __post_init__(self):
+        super().__post_init__()
+
     def forward(self, x: Circuit) -> "ndarray":
-        return super().forward(x).probabilities(self.qubits)
+        return super().forward(x).probabilities(self.qubits).reshape(1, -1)
+
+    @property
+    def output_shape(self):
+        return (1, 2 ** len(self.qubits))
 
 
 class SamplesLayer(QuantumDecodingLayer):
 
     def forward(self, x: Circuit) -> "ndarray":
         return self.backend.cast(super().forward(x).samples(), dtype=np.float64)
 
+    @property
+    def output_shape(self):
+        return (self.nshots, len(self.qubits))
+
 
 class StateLayer(QuantumDecodingLayer):
 
@@ -94,6 +105,10 @@ def forward(self, x: Circuit) -> "ndarray":
             (self.backend.np.real(state), self.backend.np.imag(state))
         )
 
+    @property
+    def output_shape(self):
+        return (2, 2**self.nqubits)
+
 
 @dataclass
 class ExpectationLayer(QuantumDecodingLayer):

src/qiboml/models/keras.py

@@ -4,6 +4,7 @@
 
 import keras
 import numpy as np
+import tensorflow as tf
 
 # from keras.src.backend import compute_output_spec
 from qibo.config import raise_error
@@ -12,9 +13,6 @@
 import qiboml.models.encoding_decoding as ed
 from qiboml.models.abstract import QuantumCircuitLayer
 
-# import tensorflow as tf
-
-
 """
 def _keras_factory(module):
     for name, layer in inspect.getmembers(module, inspect.isclass):
@@ -81,17 +79,33 @@ def __init__(self, layers: list[QuantumCircuitLayer]):
                 f"The last layer has to be a `QuantumDecodinglayer`, but is {layers[-1]}",
             )
 
-    def call(self, x: "ndarray"):
+    def call(self, x: tf.Tensor) -> tf.Tensor:
+        print(x.shape)
+        if self.backend.name != "tensorflow":
+            if self.backend.name == "pytorch":
+                self.backend.requires_grad(False)
+            x = self.backend.cast(np.array(x))
         for layer in self.layers:
             x = layer.forward(x)
+        print(x.shape)
+        if self.backend.name != "tensorflow":
+            x = tf.convert_to_tensor(np.array(x))
         return x
 
     def compute_output_shape(self):
+        return self.output_shape
+
+    @property
+    def output_shape(self):
         return self.layers[-1].output_shape
 
     @property
-    def nqubits(self):
+    def nqubits(self) -> int:
         return self.layers[0].circuit.nqubits
 
+    @property
+    def backend(self) -> "Backend":
+        return self.layers[0].backend
+
     def __hash__(self):
         return super().__hash__()

src/qiboml/models/pytorch.py

@@ -72,6 +72,12 @@ def __init__(self, layers: list[QuantumCircuitLayer]):
                     RuntimeError,
                     f"Layer \n{layer}\n is using {layer.backend} backend, but {self.backend} backend was expected.",
                 )
+        for layer in layers:
+            if len(layer.circuit.get_parameters()) > 0:
+                self.register_parameter(
+                    layer.__class__.__name__,
+                    torch.nn.Parameter(torch.as_tensor(layer.circuit.get_parameters())),
+                )
         if not isinstance(layers[-1], ed.QuantumDecodingLayer):
             raise_error(
                 RuntimeError,
@@ -101,3 +107,7 @@ def nqubits(self) -> int:
     @property
     def backend(self) -> "Backend":
         return self.layers[0].backend
+
+    @property
+    def output_shape(self):
+        return self.layers[-1].output_shape

tests/conftest.py

@@ -33,7 +33,9 @@ def get_backend(backend_name):
 def get_frontend(frontend_name):
     import qiboml
 
-    return getattr(qiboml, frontend_name)
+    frontend = getattr(qiboml, frontend_name)
+    setattr(frontend, "__str__", frontend_name)
+    return frontend
 
 
 AVAILABLE_BACKENDS = []

tests/test_models_decoding.py

@@ -0,0 +1,22 @@
+import numpy as np
+from qibo import gates
+from qibo.models import QFT
+from qibo.quantum_info import random_clifford
+
+import qiboml.models.encoding_decoding as ed
+
+
+def test_probabilities_layer(backend):
+    nqubits = 5
+    qubits = np.random.choice(range(nqubits), size=(4,), replace=False)
+    layer = ed.ProbabilitiesLayer(nqubits, qubits=qubits, backend=backend)
+    c = random_clifford(nqubits, backend=backend)
+    backend.assert_allclose(layer(c).ravel(), c().probabilities(qubits))
+
+
+def test_state_layer(backend):
+    nqubits = 5
+    layer = ed.StateLayer(nqubits, backend=backend)
+    c = random_clifford(nqubits, backend=backend)
+    real, im = layer(c)
+    backend.assert_allclose(real + 1j * im, c().state())

tests/test_models.py → tests/test_models_encoding.py

@@ -1,12 +1,7 @@
 import numpy as np
-from qibo import gates, set_backend
-from qibo.models import QFT
-from qibo.quantum_info import random_clifford
 
 import qiboml.models.encoding_decoding as ed
 
-set_backend("numpy")
-
 
 def test_binary_encoding_layer(backend):
     nqubits = 10
@@ -26,19 +21,3 @@ def test_phase_encoding_layer(backend):
     c = layer(data)
     angles = [gate.init_kwargs["theta"] for gate in c.queue if gate.name == "rz"]
     backend.assert_allclose(data, angles)
-
-
-def test_probabilities_layer(backend):
-    nqubits = 5
-    qubits = np.random.choice(range(nqubits), size=(4,), replace=False)
-    layer = ed.ProbabilitiesLayer(nqubits, qubits=qubits, backend=backend)
-    c = random_clifford(nqubits, backend=backend)
-    backend.assert_allclose(layer(c), c().probabilities(qubits))
-
-
-def test_state_layer(backend):
-    nqubits = 5
-    layer = ed.StateLayer(nqubits, backend=backend)
-    c = random_clifford(nqubits, backend=backend)
-    real, im = layer(c)
-    backend.assert_allclose(real + 1j * im, c().state())

tests/test_models_pytorch.py → tests/test_models_interfaces.py

@@ -1,15 +1,18 @@
 import inspect
 
+import keras
 import numpy as np
 import pytest
+import tensorflow as tf
 import torch
 from qibo import hamiltonians
 from qibo.config import raise_error
 from qibo.symbols import I, Z
 
 import qiboml.models.ansatze as ans
 import qiboml.models.encoding_decoding as ed
-from qiboml.models.pytorch import QuantumModel
+
+torch.set_default_dtype(torch.float64)
 
 
 def get_layers(module, layer_type=None):
@@ -34,22 +37,35 @@ def random_subset(nqubits, k):
     return np.random.choice(range(nqubits), size=(k,), replace=False).tolist()
 
 
-def output_dim(decoding_layer, input_dim, nqubits):
-    if decoding_layer is ed.ExpectationLayer:
-        return 1
-    elif decoding_layer is ed.ProbabilitiesLayer:
-        return 2**input_dim
-    elif decoding_layer is ed.SamplesLayer:
-        return input_dim
-    elif decoding_layer is ed.StateLayer:
-        return 2**nqubits
+def build_linear_layer(frontend, input_dim, output_dim):
+    if frontend.__name__ == "qiboml.models.pytorch":
+        return torch.nn.Linear(input_dim, output_dim)
+    elif frontend.__name__ == "qiboml.models.keras":
+        return keras.layers.Dense(output_dim)
+    else:
+        raise_error(RuntimeError, f"Unknown frontend {frontend}.")
+
+
+def build_sequential_model(frontend, layers):
+    if frontend.__name__ == "qiboml.models.pytorch":
+        return torch.nn.Sequential(*layers)
+    elif frontend.__name__ == "qiboml.models.keras":
+        return keras.Sequential(layers)
     else:
-        raise_error(RuntimeError, f"Layer {decoding_layer} not supported.")
+        raise_error(RuntimeError, f"Unknown frontend {frontend}.")
+
+
+def random_tensor(frontend, shape):
+    if frontend.__name__ == "qiboml.models.pytorch":
+        return torch.randn(shape)
+    elif frontend.__name__ == "qiboml.models.keras":
+        return tf.random.uniform(shape)
+    else:
+        raise_error(RuntimeError, f"Unknown frontend {frontend}.")
 
 
 @pytest.mark.parametrize("layer", ENCODING_LAYERS)
-def test_pytorch_encoding(backend, layer):
-    torch.set_default_dtype(torch.float64)
+def test_encoding(backend, frontend, layer):
     nqubits = 5
     dim = 4
     training_layer = ans.ReuploadingLayer(
@@ -59,29 +75,30 @@ def test_pytorch_encoding(backend, layer):
         nqubits, random_subset(nqubits, dim), backend=backend
     )
     encoding_layer = layer(nqubits, random_subset(nqubits, dim), backend=backend)
-    q_model = QuantumModel(
+    q_model = frontend.QuantumModel(
         layers=[
             encoding_layer,
             training_layer,
             decoding_layer,
         ]
     )
-    model = torch.nn.Sequential(
-        torch.nn.Linear(128, dim),
-        torch.nn.Hardshrink(),
-        q_model,
-        torch.nn.Linear(2**dim, 1),
+    model = build_sequential_model(
+        frontend,
+        [
+            build_linear_layer(frontend, 128, dim),
+            q_model,
+            build_linear_layer(frontend, 2**dim, 1),
+        ],
     )
-    data = torch.randn(1, 128)
+    data = random_tensor(frontend, (1, 128))
     model(data)
 
 
 @pytest.mark.parametrize("layer", DECODING_LAYERS)
 @pytest.mark.parametrize("analytic", [True, False])
-def test_pytorch_decoding(backend, layer, analytic):
+def test_decoding(backend, frontend, layer, analytic):
     if analytic and not layer is ed.ExpectationLayer:
         pytest.skip("Unused analytic argument.")
-    torch.set_default_dtype(torch.float64)
     nqubits = 5
     dim = 4
     training_layer = ans.ReuploadingLayer(
@@ -101,18 +118,20 @@ def test_pytorch_decoding(backend, layer, analytic):
         kwargs["observable"] = observable
         kwargs["analytic"] = analytic
     decoding_layer = layer(nqubits, decoding_qubits, **kwargs)
-    q_model = QuantumModel(
+    q_model = frontend.QuantumModel(
         layers=[
             encoding_layer,
             training_layer,
             decoding_layer,
         ]
     )
-    model = torch.nn.Sequential(
-        torch.nn.Linear(128, dim),
-        torch.nn.ReLU(),
-        q_model,
-        torch.nn.Linear(output_dim(layer, dim, nqubits), 128),
+    model = build_sequential_model(
+        frontend,
+        [
+            build_linear_layer(frontend, 128, dim),
+            q_model,
+            build_linear_layer(frontend, q_model.output_shape[-1], 1),
+        ],
     )
     data = torch.randn(1, 128)
     model(data)