[Bug] `qml.transforms.merge_rotations` cannot handle non-commuting observables

[Qottmann]

Related to https://github.com/PennyLaneAI/pennylane/issues/5316 Came up in https://github.com/PennyLaneAI/pennylane/pull/5396

@qml.transforms.merge_rotations
@qml.qnode(qml.device("default.qubit.legacy", wires=2, shots=50))
def circuit(order):
    qml.Permute(order, wires=(0, 1, 2))
    qml.RX(0.5, wires=0)
    qml.RX(-0.5, wires=0)
    return [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0))]

batch, fn = qml.workflow.construct_batch(circuit, level=None)((2, 1, 0))

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File ~/Xanadu/pennylane/pennylane/tape/tape.py:88, in rotations_and_diagonal_measurements(tape)
     87 try:
---> 88     rotations, diag_obs = qml.pauli.diagonalize_qwc_pauli_words(tape._obs_sharing_wires)
     89 except (TypeError, ValueError) as e:

File ~/anaconda3/envs/pennylane311/lib/python3.11/contextlib.py:81, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     80 with self._recreate_cm():
---> 81     return func(*args, **kwds)

File ~/Xanadu/pennylane/pennylane/pauli/utils.py:1142, in diagonalize_qwc_pauli_words(qwc_grouping)
   1141     if full_pauli_word[wire] != pauli_type:
-> 1142         raise ValueError("The list of Pauli words are not qubit-wise commuting")
   1143 else:

ValueError: The list of Pauli words are not qubit-wise commuting

The above exception was the direct cause of the following exception:

QuantumFunctionError                      Traceback (most recent call last)
Cell In[6], line 9
      6     qml.RX(-0.5, wires=0)
      7     return [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0))]
----> 9 batch, fn = qml.workflow.construct_batch(circuit, level=None)((2, 1, 0))

File ~/Xanadu/pennylane/pennylane/workflow/construct_batch.py:302, in construct_batch.<locals>.batch_constructor(*args, **kwargs)
    300 qnode._update_gradient_fn(tape=initial_tape)
    301 program = get_transform_program(qnode, level=level)
--> 302 return program((initial_tape,))

File ~/Xanadu/pennylane/pennylane/transforms/core/transform_program.py:509, in TransformProgram.__call__(self, tapes)
    507 if self._argnums is not None and self._argnums[i] is not None:
    508     tape.trainable_params = self._argnums[i][j]
--> 509 new_tapes, fn = transform(tape, *targs, **tkwargs)
    510 execution_tapes.extend(new_tapes)
    512 fns.append(fn)

File ~/Xanadu/pennylane/pennylane/transforms/optimization/merge_rotations.py:120, in merge_rotations(tape, atol, include_gates)
     32 r"""Quantum transform to combine rotation gates of the same type that act sequentially.
     33 
     34 If the combination of two rotation produces an angle that is close to 0,
   (...)
    117 
    118 """
    119 # Expand away adjoint ops
--> 120 expanded_tape = tape.expand(stop_at=lambda obj: not isinstance(obj, Adjoint))
    121 list_copy = expanded_tape.operations
    122 new_operations = []

File ~/Xanadu/pennylane/pennylane/tape/qscript.py:912, in QuantumScript.expand(self, depth, stop_at, expand_measurements)
    867 def expand(self, depth=1, stop_at=None, expand_measurements=False):
    868     """Expand all operations to a specific depth.
    869 
    870     Args:
   (...)
    910     RY(0.2, wires=['a'])]
    911     """
--> 912     new_script = qml.tape.tape.expand_tape(
    913         self, depth=depth, stop_at=stop_at, expand_measurements=expand_measurements
    914     )
    915     new_script._update()
    916     return new_script

File ~/Xanadu/pennylane/pennylane/tape/tape.py:181, in expand_tape(tape, depth, stop_at, expand_measurements)
    178 if tape.samples_computational_basis and len(tape.measurements) > 1:
    179     _validate_computational_basis_sampling(tape.measurements)
--> 181 diagonalizing_gates, diagonal_measurements = rotations_and_diagonal_measurements(tape)
    182 for queue, new_queue in [
    183     (tape.operations + diagonalizing_gates, new_ops),
    184     (diagonal_measurements, new_measurements),
    185 ]:
    186     for obj in queue:

File ~/Xanadu/pennylane/pennylane/tape/tape.py:99, in rotations_and_diagonal_measurements(tape)
     90     if any(isinstance(m, (ProbabilityMP, SampleMP, CountsMP)) for m in tape.measurements):
     91         raise qml.QuantumFunctionError(
     92             "Only observables that are qubit-wise commuting "
     93             "Pauli words can be returned on the same wire.\n"
   (...)
     96             "for each non-commuting observable."
     97         ) from e
---> 99     raise qml.QuantumFunctionError(_err_msg_for_some_meas_not_qwc(tape.measurements)) from e
    101 measurements = copy.copy(tape.measurements)
    103 for o, i in zip(diag_obs, tape._obs_sharing_wires_id):

QuantumFunctionError: Only observables that are qubit-wise commuting Pauli words can be returned on the same wire, some of the following measurements do not commute:
[expval(X(0)), expval(Y(0))]

[albi3ro]

A more minimal version of the problem:

@qml.transforms.merge_rotations
@qml.qnode(qml.device("default.qubit"))
def circuit():
    qml.RX(0.5, wires=0)
    qml.RX(-0.5, wires=0)
    return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0))

circuit()

Same fundamental problem as #5316 . This is due to tape.expand mixing up operator expansion with measurement validation.