Fix pytorch gradients

src/qibo/backends/numpy.py

@@ -167,7 +167,6 @@ def matrix_fused(self, fgate):
     def apply_gate(self, gate, state, nqubits):
         state = self.np.reshape(state, nqubits * (2,))
         matrix = gate.matrix(self)
-        print("matrix applied:", matrix)
         if gate.is_controlled_by:
             matrix = self.np.reshape(matrix, 2 * len(gate.target_qubits) * (2,))
             ncontrol = len(gate.control_qubits)
@@ -437,10 +436,8 @@ def execute_circuit(self, circuit, initial_state=None, nshots=1000):
                 if initial_state is None:
                     state = self.zero_state(nqubits)
 
-                print("state before circuit execution:", state)
                 for gate in circuit.queue:
                     state = gate.apply(self, state, nqubits)
-                    print("state during circuit execution:", state)
 
             if circuit.has_unitary_channel:
                 # here we necessarily have `density_matrix=True`, otherwise

src/qibo/backends/pytorch.py

@@ -121,23 +121,34 @@ def matrix_parametrized(self, gate):
         name = gate.__class__.__name__
         _matrix = getattr(self.matrices, name)
         if name == "GeneralizedRBS":
+            for parameter in ["theta", "phi"]:
+                if not isinstance(gate.init_kwargs[parameter], self.np.Tensor):
+                    gate.init_kwargs[parameter] = self.cast_parameter(
+                        gate.init_kwargs[parameter], trainable=gate.trainable
+                    )
+                elif gate.init_kwargs[parameter].requires_grad == True:
+                    gate.trainable = True
+                else:
+                    gate.trainable = False
             _matrix = _matrix(
                 qubits_in=gate.init_args[0],
                 qubits_out=gate.init_args[1],
-                theta=self.cast_parameter(
-                    gate.init_kwargs["theta"], trainable=gate.trainable
-                ),
-                phi=self.cast_parameter(
-                    gate.init_kwargs["phi"], trainable=gate.trainable
-                ),
+                theta=gate.init_kwargs["theta"],
+                phi=gate.init_kwargs["phi"],
             )
         else:
-            parameters = (
-                self.cast_parameter(param, trainable=gate.trainable)
-                for param in gate.parameters
-            )
-            _matrix = _matrix(*parameters)
-        print("parameterized matrix:", _matrix)
+            if not isinstance(gate.parameters[0], self.np.Tensor):
+                parameters = tuple(
+                    self.cast_parameter(param, trainable=gate.trainable)
+                    for param in gate.parameters
+                )
+                gate.parameters = parameters
+            elif gate.parameters[0].requires_grad == True:
+                gate.trainable = True
+                print("gate.parameters:", *gate.parameters)
+            else:
+                gate.trainable = False
+            _matrix = _matrix(*gate.parameters)
         return _matrix
 
     def cast_parameter(self, x, trainable):

tests/test_torch_gradients.py

@@ -9,21 +9,21 @@
 nepochs = 2
 optimizer = torch.optim.Adam
 target_state = torch.ones(2, dtype=torch.complex128) / 2.0
-# Define circuit ansatz
-params = torch.tensor(torch.rand(1, dtype=torch.float64), requires_grad=True)
-print(params)
+params = torch.rand(1, dtype=torch.float64, requires_grad=True)
+print("Initial params", params)
 c = models.Circuit(1)
-c.add(gates.RX(0, params[0]))
+gate = gates.RX(0, params)
+c.add(gate)
+
 optimizer = optimizer([params])
 for _ in range(nepochs):
     optimizer.zero_grad()
     c.set_parameters(params)
     final_state = c().state()
-    print("final state:", final_state)
     fidelity = torch.abs(torch.sum(torch.conj(target_state) * final_state))
     loss = 1 - fidelity
     loss.backward()
     print("loss:", loss)
     print("params.grad:", params.grad)
     optimizer.step()
-print(params)
+print("Final parameters:", params)