Enable trainable encoding layers

src/qiboml/interfaces/pytorch.py

@@ -51,6 +51,11 @@ def __post_init__(
 
         params = utils.get_params_from_circuit_structure(self.circuit_structure)
 
+        # Inform the Pytorch model that we are adding sub-modules here
+        for i, circ in enumerate(self.circuit_structure):
+            if isinstance(circ, QuantumEncoding) and circ.encoding_rule is not None:
+                self.add_module(f"enc{i}", circ.encoding_rule)
+
         params = torch.as_tensor(self.backend.to_numpy(x=params)).ravel()
         params.requires_grad = True
         self.circuit_parameters = torch.nn.Parameter(params)

src/qiboml/interfaces/utils.py

@@ -20,6 +20,11 @@ def get_params_from_circuit_structure(
     for circ in circuit_structure:
         if not isinstance(circ, QuantumEncoding):
             params.extend([p for param in circ.get_parameters() for p in param])
+        else:
+            if circ.encoding_rule is not None:
+                params.extend(
+                    [p for param in circ.circuit.get_parameters() for p in param]
+                )
     return params
 
 

src/qiboml/models/encoding.py

@@ -1,8 +1,8 @@
 import inspect
-from abc import ABC, abstractmethod
+from abc import ABC
 from dataclasses import dataclass, field
 from functools import cached_property
-from typing import Optional
+from typing import Optional, Union
 
 import numpy as np
 from qibo import Circuit, gates
@@ -19,10 +19,14 @@ class QuantumEncoding(ABC):
     Args:
         nqubits (int): total number of qubits.
         qubits (tuple[int], optional): set of qubits it acts on, by default ``range(nqubits)``.
+        encoding_rule (Optional[Union["torch.nn.module", "keras.layers.Layer"]]): optional
+            trainable encoding rule which can be used to preprocess the data with some
+            classical model.
     """
 
     nqubits: int
     qubits: Optional[tuple[int]] = None
+    encoding_rule: Optional[Union["torch.nn.module", "keras.layers.Layer"]] = None
 
     _circuit: Circuit = None
 
@@ -44,10 +48,11 @@ def _data_to_gate(self):
             f"_data_to_gate method is not implemented for encoding {self}.",
         )
 
-    @abstractmethod
     def __call__(self, x: ndarray) -> Circuit:
         """Abstract call method."""
-        pass
+        if self.encoding_rule is not None:
+            return self.encoding_rule(x)
+        return x
 
     @property
     def circuit(
@@ -107,6 +112,9 @@ def __call__(self, x: ndarray) -> Circuit:
         Returns:
             (Circuit): the constructed ``qibo.Circuit``.
         """
+        # Applying encoding rule if we have one
+        x = super().__call__(x)
+
         circuit = self.circuit
         x = x.ravel()
         for i, q in enumerate(self.qubits):
@@ -133,6 +141,9 @@ def __call__(self, x: ndarray) -> Circuit:
                 RuntimeError,
                 f"Invalid input dimension {x.shape[-1]}, but the allocated qubits are {self.qubits}.",
             )
+
+        x = super().__call__(x)
+
         circuit = self.circuit
         x = x.ravel()
         for i, q in enumerate(self.qubits):

tests/test_interfaces.py

@@ -3,6 +3,9 @@
 
 import numpy as np
 import pytest
+import tensorflow as tf
+import torch
+import torch.nn as nn
 from qibo import hamiltonians
 from qibo.config import raise_error
 from qibo.symbols import Z
@@ -13,6 +16,25 @@
 from qiboml.operations.differentiation import PSR
 
 
+class TorchLinearEncoding(nn.Module):
+    """Activation function which helps in giving more sensitivity around zero."""
+
+    def __init__(self):
+        super().__init__()
+        self.param1 = nn.Parameter(torch.tensor(np.random.randn()))
+        self.param2 = nn.Parameter(torch.tensor(np.random.randn()))
+
+    def forward(self, x):
+        return self.param1 * torch.tensor(x) + self.param2
+
+
+def construct_linear_encoding(frontend):
+    if frontend.__name__ == "qiboml.interfaces.pytorch":
+        return TorchLinearEncoding()
+    elif frontend.__name__ == "qiboml.interfaces.keras":
+        pytest.skip("Trainable encoding are supported by Pytorch interface only.")
+
+
 def get_layers(module, layer_type=None):
     layers = []
     for _, layer in inspect.getmembers(module, inspect.isclass):
@@ -282,8 +304,9 @@ def backprop_test(frontend, model, data, target):
     # specific (rare) cases
 
 
+@pytest.mark.parametrize("encoding_rule", [False, True])
 @pytest.mark.parametrize("layer,seed", zip(ENCODING_LAYERS, [6, 4]))
-def test_encoding(backend, frontend, layer, seed):
+def test_encoding(backend, frontend, encoding_rule, layer, seed):
     set_device(frontend)
     set_seed(frontend, seed)
 
@@ -308,7 +331,14 @@ def test_encoding(backend, frontend, layer, seed):
         backend=backend,
     )
 
-    encoding_layer = layer(nqubits, random_subset(nqubits, dim))
+    enc_rule = None
+    if encoding_rule:
+        enc_rule = construct_linear_encoding(frontend)
+    encoding_layer = layer(
+        nqubits=nqubits,
+        qubits=random_subset(nqubits, dim),
+        encoding_rule=enc_rule,
+    )
 
     circuit_structure = [encoding_layer, training_layer]