Circuit cutting: Add expansion to cut_circuit transform

[trbromley]

Context:

Adds an expansion function to cut_circuit, allowing support for circuits that contain WireCut operations nested in operations or sub-tapes.

[codecov]

Codecov Report

Merging #2340 (ed9aae5) into master (e8d8371) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #2340   +/-   ##
=======================================
  Coverage   99.39%   99.39%           
=======================================
  Files         233      233           
  Lines       17974    17979    +5     
=======================================
+ Hits        17865    17870    +5     
  Misses        109      109           

Impacted Files	Coverage Δ
pennylane/transforms/qcut.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update e8d8371...ed9aae5. Read the comment docs.

[anthayes92]

Thanks @trbromley, a useful addition! Just one question about how the function is defined.

Well tested, approved ✔️

[anthayes92]

I'm thrown by the indentation level here, it looks like we're defining the expand_fn method here but there some recursive subtlety here:
_qcut_expand_fn uses expand_fn which is defined by _qcut_expand_fn ...

Is this how it's working?

[josh146]

Yeah I was thrown by this as well 🤔 It makes sense now that I re-read it, but I'm curious if the following would also work?

def _qcut_expand_fn(tape, ...)
    ...
    if max_depth > 0:
        return _qcut_expand_fn.expand_fn(tape.expand(), max_depth=max_depth - 1)

cut_circuit.expand_fn = _qcut_expand_fn

This is easier for me to grok on first read

[trbromley]

Yes good question. It was originally

return _qcut_expand_fn(tape.expand(), max_depth=max_depth - 1)

on line 1236, and this works too. It just makes it hard to track with spies in the test, we'd need one spy for cut_circuit.expand_fn (first call) and another for _qcut_expand_fn (subsequent calls).

[josh146]

Looks good to me @trbromley! Only comment is on the recursion, but not a blocker to merging

[josh146]

Yeah I was thrown by this as well 🤔 It makes sense now that I re-read it, but I'm curious if the following would also work?

def _qcut_expand_fn(tape, ...)
    ...
    if max_depth > 0:
        return _qcut_expand_fn.expand_fn(tape.expand(), max_depth=max_depth - 1)

cut_circuit.expand_fn = _qcut_expand_fn

This is easier for me to grok on first read

[DSGuala]

Looks great! Thanks Tom!

I tested this with qml.TTN and it worked well, just as in the tests. This was done using the wrapper (@qml.cut_circuit) and the transform (qml.transforms.cut_circuit(circuit)).

Trying to follow the Usage Details in the qcut docs to implement manually, I had to make sure to add an expansion line:

n_wires=4
n_wires_device = 2
n_block_wires = 2
n_blocks = 2**(np.log2(n_wires/n_block_wires)+1)-1
n_SEL_layers = 2
n_SEL_wires = n_block_wires

shape = qml.StronglyEntanglingLayers.shape(n_layers=n_SEL_layers, n_wires=n_SEL_wires)
template_weights = [np.random.random(size=shape)] * int(n_blocks)
dev = qml.device('default.qubit',wires=n_wires)

def block(weights, wires):
    qml.StronglyEntanglingLayers(weights,wires)
    qml.WireCut(wires=wires)

with qml.tape.QuantumTape() as tape:
    qml.TTN(
        wires=range(n_wires),
        n_block_wires=n_block_wires,
        block=block,
        n_params_block=3,
        template_weights=template_weights,
    )
    qml.expval(qml.PauliZ(wires=n_wires-1))

###
tape = tape.expand() #this is the expansion line
###
graph = qml.transforms.qcut.tape_to_graph(tape)
qml.transforms.qcut.replace_wire_cut_nodes(graph)
fragments, communication_graph = qml.transforms.qcut.fragment_graph(graph)
fragment_tapes = [qml.transforms.qcut.graph_to_tape(f) for f in fragments]


fragment_tapes = [qml.transforms.qcut.remap_tape_wires(t, dev.wires) for t in fragment_tapes]

expanded = [qml.transforms.qcut.expand_fragment_tape(t) for t in fragment_tapes]

configurations = []
prepare_nodes = []
measure_nodes = []
for tapes, p, m in expanded:
    configurations.append(tapes)
    prepare_nodes.append(p)
    measure_nodes.append(m)

tapes = tuple(tape for c in configurations for tape in c)

results = qml.execute(tapes, dev, gradient_fn=None)

qml.transforms.qcut.qcut_processing_fn(results,communication_graph,prepare_nodes,measure_nodes,use_opt_einsum=True)

[trbromley]

Thanks @DSGuala! What is the issue you are seeing with the tape = tape.expand() line commented out? It's working ok for me with and without that line.

[DSGuala]

Thanks @DSGuala! What is the issue you are seeing with the tape = tape.expand() line commented out? It's working ok for me with and without that line.

@trbromley Sorry, I shared the wrong code. This one should give an error when we comment out that expansion line. The reason is that the tape is not getting cut into fragments.

n_wires=4
n_wires_device = 2
n_block_wires = 2
n_blocks = 2**(np.log2(n_wires/n_block_wires)+1)-1
n_SEL_layers = 2
n_SEL_wires = n_block_wires

np.random.seed(1)
shape = qml.StronglyEntanglingLayers.shape(n_layers=n_SEL_layers, n_wires=n_SEL_wires)
template_weights = [np.random.random(size=shape)] * int(n_blocks)
dev = qml.device('default.qubit',wires=n_wires_device)

def block(weights, wires):
    qml.StronglyEntanglingLayers(weights,wires)
    qml.WireCut(wires=wires)

with qml.tape.QuantumTape() as tape:
    qml.TTN(
        wires=range(n_wires),
        n_block_wires=n_block_wires,
        block=block,
        n_params_block=3,
        template_weights=template_weights,
    )
    qml.expval(qml.PauliZ(wires=n_wires-1))

###
#tape = tape.expand() #this is the expansion line
###
graph = qml.transforms.qcut.tape_to_graph(tape)
qml.transforms.qcut.replace_wire_cut_nodes(graph)
fragments, communication_graph = qml.transforms.qcut.fragment_graph(graph)
fragment_tapes = [qml.transforms.qcut.graph_to_tape(f) for f in fragments]


fragment_tapes = [qml.transforms.qcut.remap_tape_wires(t, dev.wires) for t in fragment_tapes]

expanded = [qml.transforms.qcut.expand_fragment_tape(t) for t in fragment_tapes]

configurations = []
prepare_nodes = []
measure_nodes = []
for tapes, p, m in expanded:
    configurations.append(tapes)
    prepare_nodes.append(p)
    measure_nodes.append(m)

tapes = tuple(tape for c in configurations for tape in c)
results = qml.execute(tapes, dev, gradient_fn=None)

qml.transforms.qcut.qcut_processing_fn(results,communication_graph,prepare_nodes,measure_nodes,use_opt_einsum=True)