Upgrade to Pennylane ^0.41.x

[stroblme]

PR on the journey of upgrading Pennylane (thus resolving #171 as well); Apparently there were some changes with the RZ and CRZ gates from 0.40 to 0.42 (0.41 affected as well). I extended the test model/test_ batching such that all ansatzes are being checked now.
The with ansatzes featuring RZ or CRZ gates, an error similar to cannot reshape array of size X into shape (Y,) occurs.
When changing qml.RZ to e.g. qml.RY and CRZ to CRY in our Gates class, everything runs smoothly (except of course all the other tests that actually check on the output of the model).
I tried the following already in the RZ gates class:

p = qml.math.exp(-0.5j * theta)
z = qml.math.zeros_like(p)

return qml.math.stack([stack_last([p, z]), stack_last([z, qml.math.conj(p)])], axis=-2)

signs = qml.math.array([-1, 1], like=theta)
arg = 0.5j * theta

if qml.math.ndim(arg) == 0:
    return qml.math.diag(qml.math.exp(arg * signs))

diags = qml.math.exp(qml.math.outer(arg, signs))
return diags[:, :, np.newaxis] * qml.math.cast_like(qml.math.eye(2, like=diags), diags)

and

phase = qml.math.exp(-0.5j * theta)
return qml.math.stack([phase, qml.math.conj(phase)], axis=-1)


prefactors = qml.math.array([-0.5j, 0.5j], like=theta)
if qml.math.ndim(theta) == 0:
    product = theta * prefactors
else:
    product = qml.math.outer(theta, prefactors)
return qml.math.exp(product)

in compute_matrix and compute_eigvals respectively, but without success.
Comparing to a RX or RY gate did not show any major differences, so I suspect that the reason for the different behavior is somewhere deeper.

[stroblme]

Brief update; RZ/CRZ implementation validated; the error happens around the following lines inside apply_diagonal_unitary which is called as RZ/CRZ are diagonal in z basis as checked earlier in _apply_operation_default (this does not apply to all other gates, which is why RX, RY etc. is working)

    eigvals = op.eigvals()
    eigvals = math.stack(eigvals)
    eigvals = math.reshape(eigvals, [2] * len(channel_wires))
    eigvals = math.cast_like(eigvals, state)

Because eigvals contains the batch dimension (B, EV) where EV=2 in this case, the reshape fails as (B, EV)->(2) can't work.
If no batching is applied, eigvals is just (2,) and therefore the reshape works.

[stroblme]

After some further digging; adjusting this part of the code in such a way that it works for our case will most likely break other stuff as it's a quite fundamental part of the code.
A temporary workaround that works is to comment out the following liens in _apply_operation_default in devices.qubit_mixed.apply_operation.py:

    if op in diagonal_in_z_basis:
        return apply_diagonal_unitary(op, state, is_state_batched, debugger, **_)

This will treat RZ/CRZ as standard gates und uses their matrix representation instead of going via eigenvalues.
Will create Pennylane issue and see what they think.
With this fix applied, all tests pass with the latest version of pennylane (0.42.3).

Minimal (failing) example for unpatched version:

import pennylane as qml

@qml.qnode(device=qml.device("default.mixed", wires=1))
def circuit(theta):
    qml.RZ(theta, wires=0)
    return qml.density_matrix(wires=[0])
circuit(np.array([1,2]))

with the expected result being

tensor([[[1.+0.j, 0.+0.j],
         [0.+0.j, 0.+0.j]],

        [[1.+0.j, 0.+0.j],
         [0.+0.j, 0.+0.j]]], requires_grad=True)

(2x zero state).

Thanks @majafranz for confirming this!

[stroblme]

See corresponding Pennylane Issue for further updates..

Note that while the patch suggested resolves the RZ/CRZ problematic, one still has to apply the changes suggested in this PR as well!

[stroblme]

Issue resolved by Pennylane. Let's wait for the next release ☺️