Add support for allocating qubits in decompose to cirq.unitary

cirq-core/cirq/protocols/unitary_protocol.py

@@ -179,15 +179,30 @@ def _strat_unitary_from_decompose(val: Any) -> Optional[np.ndarray]:
     if operations is None:
         return NotImplemented
 
+    all_qubits = frozenset(q for op in operations for q in op.qubits)
+    work_qubits = frozenset(qubits)
+    ancillas = tuple(sorted(q for q in all_qubits if q not in work_qubits))
+
+    ordered_qubits = ancillas + tuple(qubits)
+    val_qid_shape = (2,) * len(
+        ancillas
+    ) + val_qid_shape  # For now ancillas have only one qid_shape = (2,).
+
     # Apply sub-operations' unitary effects to an identity matrix.
     state = qis.eye_tensor(val_qid_shape, dtype=np.complex128)
     buffer = np.empty_like(state)
     result = apply_unitaries(
-        operations, qubits, ApplyUnitaryArgs(state, buffer, range(len(val_qid_shape))), None
+        operations, ordered_qubits, ApplyUnitaryArgs(state, buffer, range(len(val_qid_shape))), None
     )
 
     # Package result.
     if result is None:
         return None
+
     state_len = np.prod(val_qid_shape, dtype=np.int64)
-    return result.reshape((state_len, state_len))
+    work_state_len = np.prod(val_qid_shape[len(ancillas) :], dtype=np.int64)
+    result = result.reshape((state_len, state_len))
+    # Assuming borrowable qubits are restored to their original state and
+    # clean qubits restord to the zero state then the desired unitary is
+    # the upper left square.
+    return result[:work_state_len, :work_state_len]

cirq-core/cirq/protocols/unitary_protocol_test.py

@@ -93,6 +93,40 @@ def _decompose_(self, qubits):
         yield FullyImplemented(self.unitary_value).on(qubits[0])
 
 
+class GateThatAllocatesAQubit(cirq.Gate):
+    _target_unitary = np.array([[1, 0], [0, -1]])
+
+    def _num_qubits_(self):
+        return 1
+
+    def _decompose_(self, q):
+        anc = cirq.NamedQubit("anc")
+        yield cirq.CX(*q, anc)
+        yield (cirq.Z)(anc)
+        yield cirq.CX(*q, anc)
+
+
+class GateThatAllocatesTwoQubits(cirq.Gate):
+    # Unitary = (-j I_2) \otimes Z
+    _target_unitary = np.array([[-1j, 0, 0, 0], [0, 1j, 0, 0], [0, 0, 1j, 0], [0, 0, 0, -1j]])
+
+    def _num_qubits_(self):
+        return 2
+
+    def _decompose_(self, qs):
+        q0, q1 = qs
+        anc = cirq.NamedQubit.range(2, prefix='two_ancillas_')
+
+        yield cirq.X(anc[0])
+        yield cirq.CX(q0, anc[0])
+        yield (cirq.Y)(anc[0])
+        yield cirq.CX(q0, anc[0])
+
+        yield cirq.CX(q1, anc[1])
+        yield (cirq.Z)(anc[1])
+        yield cirq.CX(q1, anc[1])
+
+
 class DecomposableOperation(cirq.Operation):
     qubits = ()
     with_qubits = NotImplemented
@@ -201,6 +235,23 @@ def test_decompose_and_get_unitary():
     np.testing.assert_allclose(_strat_unitary_from_decompose(DummyComposite()), np.eye(1))
     np.testing.assert_allclose(_strat_unitary_from_decompose(OtherComposite()), m2)
 
+    np.testing.assert_allclose(
+        _strat_unitary_from_decompose(GateThatAllocatesAQubit()),
+        GateThatAllocatesAQubit._target_unitary,
+    )
+    np.testing.assert_allclose(
+        _strat_unitary_from_decompose(GateThatAllocatesAQubit().on(a)),
+        GateThatAllocatesAQubit._target_unitary,
+    )
+    np.testing.assert_allclose(
+        _strat_unitary_from_decompose(GateThatAllocatesTwoQubits()),
+        GateThatAllocatesTwoQubits._target_unitary,
+    )
+    np.testing.assert_allclose(
+        _strat_unitary_from_decompose(GateThatAllocatesTwoQubits().on(a, b)),
+        GateThatAllocatesTwoQubits._target_unitary,
+    )
+
 
 def test_decomposed_has_unitary():
     # Gates