split_non_commuting transform for expectation values

[Qottmann]

Context:

Currently, qml.expval() does not allow for lists of operators that are not commuting. I.e. the following code

import pennylane as qml
with qml.tape.QuantumTape() as tape:
    qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
    qml.expval(qml.PauliY(0))
qml.tape.tape.expand_tape(tape)

produces an error message as the observables do not commute.

Description of the Change:

The idea is to split the tape into groups of commuting observables, compute those, and recombine them.
We want to do this in the device class on the call of batch_transform() in a similar way as Hamiltonians are expanded.

For this we first need the list of observables, which in the current state of the operator and observable classes is a bit cumbersome. This should be easier once those classes are updated.
Now this list of observables is grouped into commuting groups using qml.groupings.group_observables(). We keep track of the indices by also passing a list of dummy coefficients.

Then a new tape for each of those groups is generated and passed. In the end we recombine all results into a single list with the same order as the observables have been initially given.

Re-structuring after grouping

We want the user not to know that this grouping has taken place internally and output the same shape and order as the input of expectation values. For this we need to pass a non-trivial classical processing function. I decided to go with the following:

def reorder_fn(res):
    return qml.math.concatenate(res)[qml.math.concatenate(group_coeffs)]

res is a list of results for each commuting group. So the first step, concatenating those results produces a list of the original size. The second step is re-ordering the results according to the original input. An example for the following circuit

import pennylane as qml
dev = qml.device("default.qubit", wires=5)
@qml.qnode(dev)
def circ():
    return [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
    qml.expval(qml.PauliY(0)),
    qml.expval(qml.PauliZ(1)),
    qml.expval(qml.PauliX(1) @ qml.PauliX(4)),
    qml.expval(qml.PauliX(3))]

Here the we would internally have res = [tensor([1., 1., 0.], requires_grad=True), tensor([0., 0.], requires_grad=True)] and group_coeffs = [array([0, 2, 4]), array([1, 3])]. So reorder_fn computes

reorder_fn(res) = tensor([1., 1., 0., 0., 0.,], requires_grad=True)[array(0,2,4,1,3)] = tensor([1., 0., 0., 1., 0.], requires_grad=True)

Status at initial WIP PR

With the current changes, cases like the following work:

import pennylane as qml
dev = qml.device("default.qubit", wires=5)
@qml.qnode(dev)
def circ():
    return [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
    qml.expval(qml.PauliY(0)),
    qml.expval(qml.PauliZ(1)),
    qml.expval(qml.PauliX(1) @ qml.PauliX(4)),
    qml.expval(qml.PauliX(3))]
circ()

import pennylane as qml
dev = qml.device("default.qubit", wires=2)
@qml.qnode(dev)
def circ():
    qml.QubitStateVector([0.+0.j, 0.+0.j, 0.+1.j, 0.+0.j], wires=range(2))
    return [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
    qml.expval(qml.PauliY(0)),
    qml.expval(qml.PauliZ(1))]
circ()

import pennylane as qml
dev = qml.device("default.qubit", wires=2)
@qml.qnode(dev)
def circ():
    qml.PauliY(0)
    return [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
    qml.expval(qml.PauliY(0)),
    qml.expval(qml.PauliZ(1))]
circ()

import pennylane as qml
dev = qml.device("default.qubit", wires=2)
H = qml.Hamiltonian([0.5,0.5],[qml.PauliZ(0), qml.PauliZ(1)])
@qml.qnode(dev)
def circ():
    qml.PauliY(0)
    return qml.expval(H), qml.expval(H)
circ()

However, some of the tests in the test suite still fail, so I need to figure out why that is and make changes accordingly.

Possible Drawbacks:

Statevector devies can compute non-commuting observables straight away so they should be treated as a special case for speed ups in a separate PR.

Related GitHub Issues:

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

Merging #2587 (0fcc85d) into master (4dda6cf) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #2587   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         243      244    +1     
  Lines       19633    19667   +34     
=======================================
+ Hits        19553    19587   +34     
  Misses         80       80           

Impacted Files	Coverage Δ
pennylane/_device.py	98.08% <100.00%> (+0.01%)	⬆️
pennylane/transforms/__init__.py	100.00% <100.00%> (ø)
pennylane/transforms/split_non_commuting.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 4dda6cf...0fcc85d. Read the comment docs.

[Qottmann]

I just realized that the failure from the local tests that I ran are likely due to some dependency problems in my pennylane environment as the tests seem to pass here.

I will then continue and add new tests with non-commuting lists of observables.

[albi3ro]

I'd recommend moving all of this into its own function and then calling that function here.
How about in a new file pennylane/transforms/split_noncommuting_observables.py or something like that. The name needs a little work.

Devices are already highly bloated classes responsible for too many different things. What we want to do is create a self-contained black box that solves a problem. We abstract away the implementation details from places that don't need to know the implementation details.

This will improve readability in _device.py and allow us to repurpose the same function later on as other things change. When we rewrite how we do devices, we'd still like to be able to use the same function, even if we use it in a different place.

If you notice in your chunk of code, it doesn't call any other device methods. That's a good sign it can be extracted somewhere else.

Small independent functions are also much easier to unit test.

[Qottmann]

Noticed a bug: reorder_fn(res) did not return the right ordering. Now with a hacky solution, maybe this can be improved:

def reorder_fn(res):
    """re-order the output to the original shape and order"""
    res = qml.math.concatenate(res)
    new_res = res.copy() # to keep the same format as res
    reorder_indxs = qml.math.concatenate(group_coeffs)
    for i,out in zip(reorder_indxs, res):
        new_res[i] = out

Further, this does not work with a mixture of Hamiltonians and regular observables. Since everything that is multiplied by a number becomes a Hamiltonian, this might be a bit restrictive in use cases for this addition.

[Qottmann]

Thank you @albi3ro for the feedback, I now moved everything to qml.transforms.split_non_commuting(). Since this function is already tested via test_expval_non_commuting_observables() in test_measurements.py, is there any benefit to adding more tests like in tests/transform/test_hamiltonian_expand.py? If yes, what kind of tests would further be required?

[josh146]

Nice work @Qottmann! I've given it a quick first pass :)

[Qottmann]

Ready for review @josh146 @albi3ro

Main update is that now it includes tests for different interfaces (jax-jit cannot handle multiple return values in qfuncs atm).

Do you know how to get rid of the codefactor error pennylane/_device.py#L732 Access to a protected member _obs_sharing_wires of a client class (protected-access) I saw this attribute being used in other parts of the codebase so I dont understand why it is a problem here?

[josh146]

Thanks @Qottmann! This is looking very close, all my suggestions were minor. I only had two major ones:

• (Optional) turn this into a batch transform, so people can use it on QNodes :)

• Double check the gradient tests regarding the zero Jacobians 🤔

[Qottmann]

Even though I moved the entry in the doc string of pennylane/transforms/__init__.py to "Transforms that act on QNodes", in the deployed docs it is still listed under "Transforms that act on tapes".

I am not 100% sure how the doc-string should be written now that it is a batch_transform. I tried to cover both in the input and return descriptions, but focus on QNodes in the example and put the tape examples into usage details.

Other than that I think it is good to go.

[josh146]

Great work @Qottmann! I read through the docstring, and it is very impressive - you did a good job describing both the functional (QNode) and object oriented (tape) user-interface 💪

[antalszava]

[sc-17239]