feat: added test for moments with callback gates

qiboteam · Mar 11, 2024 · 39ba5d4 · 39ba5d4
1 parent 19e884c
commit 39ba5d4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/qibo/models/circuit.py b/src/qibo/models/circuit.py
@@ -94,7 +94,7 @@ def moments(self):
         moments = [self.nqubits * [None]]
         moment_index = self.nqubits * [0]
         for gate in self:
-            if gate.qubits:
+            if not isinstance(gate, gates.CallbackGate):
                 qubits = gate.qubits
             else:  # special gate acting on all qubits
                 qubits = tuple(range(self.nqubits))

diff --git a/tests/test_models_circuit.py b/tests/test_models_circuit.py
@@ -19,8 +19,10 @@ def test_parametrizedgates_class():
 
 
 def test_queue_class():
+    from qibo.callbacks import EntanglementEntropy
     from qibo.models.circuit import _Queue
 
+    entropy = EntanglementEntropy([0])
     queue = _Queue(4)
     gatelist = [
         gates.H(0),
@@ -29,12 +31,14 @@ def test_queue_class():
         gates.H(2),
         gates.CNOT(1, 2),
         gates.Y(3),
+        gates.CallbackGate(entropy),
     ]
     for g in gatelist:
         queue.append(g)
     assert queue.moments == [
         [gatelist[0], gatelist[1], gatelist[3], gatelist[5]],
         [gatelist[2], gatelist[4], gatelist[4], None],
+        [gatelist[6] for _ in range(4)],
     ]