Support symbolic decomposition for state prep via alias sampling

[tanujkhattar]

Fixes https://github.com/quantumlib/Qualtran/issues/951

In general, the cirq-interop decomposition assumes access to individual bits. So anytime we would do bb.split / bb.join in the decomposition, its fine to use the cirq style decomposition as well but if we are not doing any splits or joins (example: for higher level bloqs like state prep alias sampling) its useful to migrate the cirq style decomposition to bloq style decomposition so we can support decomposing bloqs with symbolic register sizes.

More such migrations would follow as we discover bloqs which currently have cirq-style decompose_from_registers implemented but can benefit from a symbolic decomposition. An alternate strategy would be to update the cirq-interop code to support symbolics wherever we can.

[tanujkhattar]

Something is going wrong with the cirq interop again due to which cirq simulation tests for reflection using prepare are failing. This is probably related to the few other places we've seen this happen recently. I'll have to spend more time to figure out exactly whats going wrong.

[anurudhp]

Possible cause: GreedyQubitManager, even with maximize_reuse=True, does not reuse qubits in a bloq style decomposition. For example see the outputs below (for a decomp. with two chained allocate/free pairs):

Script

from typing import Iterator

import cirq
from attrs import frozen
from numpy.typing import NDArray
from qualtran import Bloq, BloqBuilder, GateWithRegisters, Signature, SoquetT
from qualtran.bloqs.basic_gates import CNOT
from qualtran.cirq_interop.testing import GateHelper

@frozen
class MultiAlloc(Bloq):
    @property
    def signature(self) -> "Signature":
        return Signature.build(q=1)

    def build_composite_bloq(
        self, bb: "BloqBuilder", q: "SoquetT"
    ) -> dict[str, "SoquetT"]:
        for _ in range(2):
            a = bb.allocate(1)
            a, q = bb.add(CNOT(), ctrl=a, target=q)
            bb.free(a)

        return {"q": q}

class MultiAllocGWR(GateWithRegisters):
    @property
    def signature(self) -> "Signature":
        return Signature.build(q=1)

    def decompose_from_registers(
        self, *, context: cirq.DecompositionContext, **quregs: NDArray[cirq.Qid]
    ) -> Iterator[cirq.OP_TREE]:
        (q,) = quregs["q"]
        for _ in range(2):
            (a,) = context.qubit_manager.qalloc(1)
            yield CNOT().on(a, q)
            context.qubit_manager.qfree([a])

def visualize(bloq: Bloq | GateWithRegisters):
    print(bloq)
    gate = GateHelper(
        gate=bloq,
        context=cirq.DecompositionContext(
            cirq.GreedyQubitManager(prefix="_greedy", maximize_reuse=True)
        ),
    )
    print(gate.circuit)
    print()
    print(gate.decomposed_circuit)
    print()

if __name__ == "__main__":
    visualize(MultiAlloc())
    visualize(MultiAllocGWR())

output

MultiAlloc
q: ───q───

_greedy_0: ───alloc────@───free─────────────────────────
                       │
_greedy_1: ───alloc────┼────────────────────@───free────
                       │                    │
q: ───────────Y^-0.5───@───Y^0.5───Y^-0.5───@───Y^0.5───

MultiAllocGWR
q: ───MultiAllocGWR───

_greedy_0: ───@───@───
              │   │
q: ───────────X───X───

[tanujkhattar]

Yup, this is known. See https://github.com/quantumlib/Qualtran/pull/963#issue-2301621495

[tanujkhattar]

@mpharrigan This is ready for a review

[mpharrigan]

@tanujkhattar do you want to merge this

[tanujkhattar]

I've addressed the nits and also added an assertion to post init to verify that we only do state prep over 1D register. Opened https://github.com/quantumlib/Qualtran/issues/1280 to track improvements

Merging now.