Shape error for tensor-like parameters in pauli arithmetic fixed (Addition to PR: #6562)

[ldi18]

Context:
I added Pauli string decompositions for single qubit parametrized gates in PR: #6562. Pennylane exploits these Pauli string representations when, for example, computing

(qml.RX(1.11, wires=0) @ qml.RY(1.23, wires=0)).matrix()

The example above works as intended, but

(qml.RX([1.11, 2.11, 3.11], wires=0) @ qml.RY([1.23, 2.23, 3.23], wires=0)).matrix()     # shape error

or

(qml.RX(np.array([1.11, 2.11, 3.11]), wires=0) @ qml.RY(np.array([1.23, 2.23, 3.23]), wires=0)).matrix()     # shape error

returns a shape error. The code above should work since RXand RYsupport tensor-like (flattened) parameters as it can be seen from

qml.RX([1.11, 2.11, 3.11], wires=0).matrix()    # works!

Description of the Change:

I added an outer product in pauli_arithmetic.

Benefits:

• Pauli representations can now be leveraged for the matrix conversion of products of ops.
• A Pauli representation of the PhaseShift gate could be added - this was not possible before as the PCPhase gate used the pauli_arithmetic, which gave an error in the pytests.

Possible Drawbacks:
None

Related GitHub Issues:
#6243 and PR: #6562

[ldi18]

@dwierichs Should I close this PR? Without the pauli reps for parametrized gates in #6562 the Pauli arithmetic won't be used so there is no need to enable batched computation - or should the bug be corrected in case pauli reps for parametrized gates are added at a later stage?

[dwierichs]

Thanks @ldi18 ,
we're still figuring out how to proceed with this one, thanks for your patience!

[dwierichs]

Hi @ldi18
The plan is to move ahead with this bug fix. Thank you for contributing it and for your patience.
Can I ask you to update it to the reduced version of the other PR? :)
@astralcai will provide a review afterwards.

[ldi18]

Hi @dwierichs and @astralcai,

this is the updated PR but tests are still missing. The tests I created previously can't be used directly anymore as we didn't include the parametrized ops in PR #6562. We could overwrite some parametrized ops with

class RX_with_pauli_rep(RX):
    @property
    def pauli_rep(self):
        if self._pauli_rep is None:
            self._pauli_rep = qml.pauli.PauliSentence(
                {
                    qml.pauli.PauliWord({self.wires[0]: "I"}): qml.math.cos(self.data[0] / 2),
                    qml.pauli.PauliWord({self.wires[0]: "X"}): -1j * qml.math.sin(self.data[0] / 2),
                }
            )
        return self._pauli_rep

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

class RY_with_pauli_rep(RY):
    @property
    def pauli_rep(self):
        if self._pauli_rep is None:
            self._pauli_rep = qml.pauli.PauliSentence(
                {
                    qml.pauli.PauliWord({self.wires[0]: "I"}): qml.math.cos(self.data[0] / 2),
                    qml.pauli.PauliWord({self.wires[0]: "Y"}): -1j * qml.math.sin(self.data[0] / 2),
                }
            )
        return self._pauli_rep

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

in the test files. Then

(RX_with_pauli_rep([1.5, 1, 1], wires=0) @ RY_with_pauli_rep([1.5, 1, 1], wires=0)).matrix()

gives the expected shape error - and no error once the changes in pauli_arithmetic proposed in this PR are included.

[dwierichs]

Thanks @ldi18
As you already have that test with modified rotation gates, I think it would be great to include those!
Maybe in addition, could you add a (very minimal) test for the particular methods that are being modified? So, for example, something like

@pytest.mark.parametrize("coeff0", [np.array([0.6, 0.2, 4.3]), -0.7])
@pytest.mark.parametrize("coeff1", [np.array([1.2, -0.9, 2.7]), -0.7])
def test_to_mat_with_broadcasting(coeff0, coeff1):
    wire_order = [0, 1, "a", "b"]
    pw0 = qml.pauli.PauliWord({0: "X", "a": "Y"})
    pw1 = qml.pauli.PauliWord({0: "Z", 1: "Y", "b": "Y"})
    ps = qml.pauli.PauliSentence({pw1: coeff0, pw1: coeff1})
    mat0 = ps.to_mat(wire_order=wire_order)
    mat1 = qml.matrix(ps.operation(), wire_order=wire_order)
    assert np.allclose(mat0, mat1)

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.62%. Comparing base (9d9205a) to head (9bf6366).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6587   +/-   ##
=======================================
  Coverage   99.62%   99.62%           
=======================================
  Files         468      468           
  Lines       44144    44144           
=======================================
  Hits        43977    43977           
  Misses        167      167           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[ldi18]

Thanks @ldi18 As you already have that test with modified rotation gates, I think it would be great to include those! Maybe in addition, could you add a (very minimal) test for the particular methods that are being modified? So, for example, something like

@pytest.mark.parametrize("coeff0", [np.array([0.6, 0.2, 4.3]), -0.7])
@pytest.mark.parametrize("coeff1", [np.array([1.2, -0.9, 2.7]), -0.7])
def test_to_mat_with_broadcasting(coeff0, coeff1):
    wire_order = [0, 1, "a", "b"]
    pw0 = qml.pauli.PauliWord({0: "X", "a": "Y"})
    pw1 = qml.pauli.PauliWord({0: "Z", 1: "Y", "b": "Y"})
    ps = qml.pauli.PauliSentence({pw1: coeff0, pw1: coeff1})
    mat0 = ps.to_mat(wire_order=wire_order)
    mat1 = qml.matrix(ps.operation(), wire_order=wire_order)
    assert np.allclose(mat0, mat1)

@dwierichs That's a great suggestion and seems sufficient to test my changes. I didn't add my tests because they are very hypothetical and don't test the changes as directly your suggestion. Using arrays as input actually caused an ambiguous truth value in the ps.operation() call

     def operation(self, wire_order=None):
        """Returns a native PennyLane :class:`~pennylane.operation.Operation` representing the PauliSentence."""
        if len(self) == 0:
            return qml.s_prod(0, Identity(wires=wire_order))
    
        summands = []
        wire_order = wire_order or self.wires
        for pw, coeff in self.items():
            pw_op = pw.operation(wire_order=list(wire_order))
            rep = PauliSentence({pw: coeff})
>           summands.append(pw_op if coeff == 1 else SProd(coeff, pw_op, _pauli_rep=rep))
E     ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I implemented the tests with lists instead ([0.1, 0.1, 1] == 1 is False and not ambiguous).

Probably the error for array input would surface again at some point, so I also added qml.math.all(coeff == 1) and a corresponding test to this PR. If you prefer to leave it out of this PR I can remove it of course.

[astralcai]

LGTM! 🚀

[dwierichs]

I'd second what @astralcai said about the changelog entry, and I like the additional fix in .operation() 💯

Thanks for your contribution @ldi18 🎉