Add support for allocating qubits in decompose to cirq.unitary

[NoureldinYosri]

This is the simplest implementation to do this in order to unblock cirq-ft.
This implementation only satisfies the first requirement of #6101 (comment)
In followup PRs I will address other requirements.

[tanujkhattar]

Can you add a test where a GateA decomposes into a GateB where GateB allocates new qubits and acts on

a) <= 4 qubits
b) > 4 qubits ?

My guess is that Case-(b) above would fail the test right now. The reason is as follows:

1. _try_decompose_into_operations_and_qubits internally calls cirq.decompose_once for gates / operations etc.

2. cirq.unitary uses (1) to get a list of decomposed operations and then delegates the logic of applying the unitary effect of these decompositions to apply_unitaries protocol.
   

   Cirq/cirq-core/cirq/protocols/unitary_protocol.py

   Line 185 in 36d67c1

   result = apply_unitaries(

3. apply_unitary then acts on each operation in the first-level decomposition and then tries to "apply" it's unitary effect onto the given target tensor. To decide a strategy, it looks at len(args.axes).
   

   Cirq/cirq-core/cirq/protocols/apply_unitary_protocol.py

   Lines 346 to 357 in 36d67c1

   	if len(args.axes) <= 4:
   	strats = [
   	_strat_apply_unitary_from_apply_unitary,
   	_strat_apply_unitary_from_unitary,
   	_strat_apply_unitary_from_decompose,
   	]
   	else:
   	strats = [
   	_strat_apply_unitary_from_apply_unitary,
   	_strat_apply_unitary_from_decompose,
   	_strat_apply_unitary_from_unitary,
   	]

4. Thus, if the unitary operators on <= 4 qubits, the recursive strategy will delegate to cirq.unitary(OpForGateB) to get its (narrow) unitary and thus be able to apply the effect of the unitary on the target tensor correctly.

5. But, if it acts > 4 qubits, then it'll try the _strat_apply_unitary_from_decompose which will fail because it currently assumes that decomposition returns an op-tree that's acts only the returned set of qubits and thus we'd need to make similar changes to _strat_apply_unitary_from_decompose as well.

I think a better strategy would be to change the return type of _try_decompose_into_operations_and_qubits to return work_qubits and ancilla_qubits instead of only qubits. But that would be a bigger change since it would trigger every call site that makes this assumption. We can leave a comment on the issue and track this as an independent improvements (I can also take care of it as part of restructuring the cirq.decompose protocol).

[NoureldinYosri]

it doesn't matter whether the gate acts on $\leq 4$ or $&gt; 4$ qubits as long as the gate doesn't implement _apply_unitary_ or _unitary_.

the issue isn't _try_decompose_into_operations_and_qubits it's that the apply_unitary protocol makes the same assumption as the unitary protocol. that the gates/operations won't allocate new qubits if it decomposes them.

I'm working a fix for the apply unitary protocol, will update this PR soon.

[tanujkhattar]

the issue isn't _try_decompose_into_operations_and_qubits it's that the apply_unitary protocol makes the same assumption as the unitary protocol

Well, what I meant was that most callsites of _try_decompose_into_operations_and_qubits would make a similar assumption; for example the strat_act_on_from_apply_decompose in simulation state:

Cirq/cirq-core/cirq/sim/simulation_state.py

Line 296 in ebec38b

operations, qubits1, _ = _try_decompose_into_operations_and_qubits(val)

Thus, changing the return signature of this method would be a good way to find and fix all such call-sites that assume that gates / operations cannot allocate new qubits as part of their decomposition.

Although, I quick grep search shows the following call sites:

cirq-core/cirq/sim/simulation_state.py:from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
cirq-core/cirq/sim/simulation_state.py:    operations, qubits1, _ = _try_decompose_into_operations_and_qubits(val)
cirq-core/cirq/protocols/kraus_protocol.py:from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
cirq-core/cirq/protocols/kraus_protocol.py:        operations, _, _ = _try_decompose_into_operations_and_qubits(val)
cirq-core/cirq/protocols/decompose_protocol.py:def _try_decompose_into_operations_and_qubits(
cirq-core/cirq/protocols/unitary_protocol.py:from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
cirq-core/cirq/protocols/unitary_protocol.py:    operations, qubits, val_qid_shape = _try_decompose_into_operations_and_qubits(val)
cirq-core/cirq/protocols/has_unitary_protocol.py:from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
cirq-core/cirq/protocols/has_unitary_protocol.py:    operations, _, _ = _try_decompose_into_operations_and_qubits(val)
cirq-core/cirq/protocols/apply_unitary_protocol.py:from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
cirq-core/cirq/protocols/apply_unitary_protocol.py:    operations, qubits, _ = _try_decompose_into_operations_and_qubits(val)
cirq-core/cirq/protocols/mixture_protocol.py:from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
cirq-core/cirq/protocols/mixture_protocol.py:        operations, _, _ = _try_decompose_into_operations_and_qubits(val)

Out of these, the has_* protocols ignore qubits and only care about decomposed operations; so they should be fine. We are fixing unitary and apply_unitary as part of this PR and simulation state should be fixed as part of #6108 (cc @senecameeks)

The kraus, mixture, apply_channel and apply_mixture protocols don't yet support a _strat_from_decompose mode and thus only work if _kraus_, _mixture_, _apply_channel_ or _apply_mixture_ protocols are implemented on val correspondingly. Therefore, for the sake of this issue, we can ignore these protocols.

[NoureldinYosri]

@tanujkhattar got it. however fixing different places where it's called depends on the logic executed there. that's why the method is not the issue for this PR since we already know that these two protocols don't support ancillas.

The apply_unitaries method is more challenging than the cirq.unitary protocol. the reason is that it can recurse/decompose multiple times each time potentially creating new ancillas which won't have place in the buffer allocated in the first call.

cirq.unitary is essentially just convenience method for handling different cases with a fall back to apply_unitaries so fixing this method is key.

To that end I have two ideas. but I don't like either of them.

[tanujkhattar]

@NoureldinYosri I think you should also look at the approach suggested here: #6040 (comment) and implemented as part of #6081. The two issues seem related.

[NoureldinYosri]

@tanujkhattar the issue for me wasn't keeping track of the new qubits. I have a couple of ways for doing that. The issue is how not to sacrifice the efficiency of the protocol.

The current implementation is not efficient in any way. but it will do the trick & unblock cirq-ft. later I'll improve its efficiency.

btw currently there is no mention of the qubit manager. when we introduce the qubit manager abstraction into cirq we will need to add it here. There is a natural way of doing this by adding the qubit manager as a property of the args.

For cirq-ft this is currently not a problem since all decompositions use the global qubit manager and there is no object to pass around.

what do you think?

[tanujkhattar]

Looking good overall. Main suggestion is to extend the ApplyUnitaryArgs class to provide methods to allocate new space for the ancillas and extract out the unitary corresponding to working qubits.

[tanujkhattar]

I think we should create a kronecker_product and factor method on ApplyUnitaryArgs; similar to SimulationState; and then dispatch the logic of allocate new space and extracting out the unitary corresponding to a subspace using function calls to these methods.

[NoureldinYosri]

I addressed the comment we had offline.

I'll keep the code that factors out the unitary as it's for now and I'll add a general factor_unitary in a followup PR to cirq-core/cirq/linalg/transformations.py

[tanujkhattar]

LGTM % minor refactoring nits

[tanujkhattar]

One comment on simple_gates.py -- either we don't create a new file in this PR or we add tests to make sure the new file has coverage.

[NoureldinYosri]

@tanujkhattar PTAL. I added sample_gates_test.py and addressed the other comments.

[tanujkhattar]

Here is a commit on top of this PR which suggests a few changes (described below) - tanujkhattar@14dd571

The tests for sample_gates.py right now test a) the target given unitary as hardcoded in the original class and b) the reduced unitary assuming the cirq.unitary implementation is correct. Ideally, we should the tests should verify that the decomposition is correct (i.e. have an assertion on the structure of the decomposed unitary).

I've refactored the sample_gates.py file so that it now contains two gates - PhaseUsingCleanAncilla and PhaseUsingDirtyAncilla each of which apply a phase to a given basis state using clean / dirty ancillas.

I've also fixed a few styling nits like added docstrings and modified function names. Can you please incorporate the changes from my commit into your PR and then we can go ahead and merge?

[tanujkhattar]

LGTM! Can you please fix the coverage failure so we can merge?

Thanks for all the help!

[NoureldinYosri]

done. one thing I noticed is that the Isolated pytest Ubuntu test is flaky, it sometimes fails on unrelated tests but rerunning it makes it succeed.

[shef4]

Hi, I set up my dev enviroment to start working on this issue and read through the the previous conversation and code.

I'm gonna work on the 2nd question you raised #6101

2nd Question: Can it support the situtation where some of the operations in the decomposition are not unitary but the overall effect on the work qubits (not the ancillas) is unitary (e.g. erasing entanglement using measurement)?

[NoureldinYosri]

@shef4 sounds good, thanks for your help

In the mean time I will soon add the consistancy and correctness checks.

[shef4]

@NoureldinYosri sure, thing! I looked into the problem more and wanted to get some input before coding.

From what I understood from the comments and the code files. I think it makes sense to add a simple test for "earsing entagled qubits using measurment" in apply_unitary_protocol_test.py that utilizes the apply_unitaries function.

The apply_unitaries function currently terminates if one of the the given unitary values does not have a unitary effect ,but does not rollback changes from previous unitray values.

A bruteforce approach could be to apply all unitary operations and check at the end if the work qubit is still unitary. This would probaly require the least amount of change to the code, but would make apply_unitaries less effecient. I could also add a separate function.

Just to clearify, can the work qubit be non-unitary during the operation sequence as long as it is unitary at the end?

[NoureldinYosri]

A bruteforce approach could be to apply all unitary operations and check at the end if the work qubit is still unitary.

@shef4 the product of unitaries is always a unitary.

the problem we have is different. if some operation is not unitary and in that case there is no _apply_unitary_ or _unitary_ methods by definition. however multiple non unitary operations could lead to a unitary operation. take for example the CCNOT gate which flips the last qubit if and only if the first 2 are ones. its unitary is given by

$$ \begin{bmatrix} 1 & 0 & 0 & 0 & 0 & 0 & 0 & 0\\ 0 & 1 & 0 & 0 & 0 & 0 & 0 & 0\\ 0 & 0 & 1 & 0 & 0 & 0 & 0 & 0\\ 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0\\ 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0\\ 0 & 0 & 0 & 0 & 0 & 1 & 0 & 0\\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 1\\ 0 & 0 & 0 & 0 & 0 & 0 & 1 & 0\\ \end{bmatrix} $$

This gate can be implemented using And gate as

class CCNOTUsingAnd(cirq_ft.GateWithRegisters):

    @cached_property   
    def registers(self):
        return cirq_ft.Registers.build(control=2, target=1)

    def decompose_from_registers(
        self, *, context: cirq.DecompositionContext, **quregs: Sequence[cirq.Qid]
    ) -> cirq.OP_TREE:
        control, target = quregs['control'], quregs['target']
        (ancilla,) = context.qubit_manager.qalloc(1)

        yield cirq_ft.And().on(*control, ancilla)
        yield cirq.CNOT(ancilla, t)
        yield cirq_ft.And(adjoint=True)(c0, c1, ancilla)

so task 2 is to support something like

u = cirq.unitary(CCNOTUsingAnd)   # Must succeeds to build the unitary.
assert np.allclose(u, cirq.unitary(cirq.CCNOT))   # Must be the same as the unitary of a CCNOT

currently the line cirq.unitary(CCNOTUsingAnd) would fail since And(adjoint=True) is not a unitary operation since its decomposition contains a measurement . the purpose of task 2 is support this and simillar use cases.

[shef4]

@NoureldinYosri Thanks for the clarification and code example, there was a typo in my brute-force solution comment.

what I meant was that from my understanding, cirq.unitary(CCNOTUsingAnd) call's _unitary_from_decompose() under the hood to evaluate CCNOTUsingAnd.

_unitary_from_decompose() then calls apply_unitaries() to produce a unitary matrix from the CCNOTUsingAnd.

I think i know where to go from here, thank you again. I'm gonna write a simple test for this example (cirq.unitary(CCNOTUsingAnd)) and watch where it fails.

[NoureldinYosri]

@shef4 please check #6182