Source code for qailab.torch.qlayer

 1""" Module with QLayer """
 2import math
 3
 4import torch
 5from torch import Tensor, nn
 6
 7from qiskit import QuantumCircuit, transpile
 8from qiskit.circuit.library.generalized_gates.isometry import Isometry
 9
10from quantum_launcher import QuantumLauncher
11from quantum_launcher.routines.qiskit_routines import QiskitBackend
12
13from qailab.circuit.utils import filter_params
14from qailab.qlauncher import CircuitProblem, ForwardPass, BackwardPass
15from qailab.torch.autograd import ExpVQCFunction
16Isometry.__init__.__defaults__ = (1e-6,)  # FIXME: If anyone has any idea, feel free
17
18

[docs]
19class QLayer(nn.Module):
20    """
21    Base quantum layer class.
22
23    This layer will output 2^num_measured_qubits features,
24    which represent the distribution of measurements from the underlying quantum circuit.
25
26    The bitstring order is [0, 1, 2, ..., 2^num_measured_qubits-1]
27    """
28    theta_trainable: Tensor
29
30    def __init__(
31        self,
32        circuit: QuantumCircuit,
33        *,
34        backend: QiskitBackend | None = None,
35        shots: int = 1024,
36    ) -> None:
37        super().__init__()
38
39        self.theta_trainable = nn.Parameter(
40            torch.empty((len(filter_params(circuit, 'weight')), 1))
41        )
42
43        self.reset_parameters()
44
45        if backend is None:
46            backend = QiskitBackend('local_simulator')
47
48        self.circuit = transpile(circuit, backend.sampler.backend) if hasattr(backend.sampler, 'backend') else circuit
49        self.circuit_pr = CircuitProblem(self.circuit)
50
51        # Helper variables for hybrid networks
52        self.in_features = len(filter_params(circuit, 'input'))
53        self.out_features = 2**self.circuit.num_clbits
54
55        self.launcher_forward = QuantumLauncher(
56            self.circuit_pr,
57            ForwardPass(shots=shots),
58            backend
59        )
60        self.launcher_backward = QuantumLauncher(
61            self.circuit_pr,
62            BackwardPass('param_shift', shots=shots),
63            backend
64        )
65

[docs]
66    def reset_parameters(self) -> None:
67        """ Parameter reset """
68        nn.init.uniform_(self.theta_trainable, 0, 2 * math.pi)

69

[docs]
70    def extra_repr(self) -> str:
71        return f"{self.circuit}"

72

[docs]
73    def forward(self, input_tensor: Tensor) -> Tensor:
74        """Forward run"""
75        out = ExpVQCFunction.apply(
76            input_tensor,
77            self.theta_trainable[:, 0],
78            self.launcher_forward,
79            self.launcher_backward
80        )
81        if not isinstance(out, torch.Tensor):
82            raise ValueError("Function did not return tensor output")
83        return out