[BUG] mpl drawer does not show expanded circuit when used with expansion_strategy='device'

[qcabepsilon]

Expected behavior

I simulate a VQE of the H2 molecule using the UCCSD ansatz from pennylane qchem with the 'default.qubit' device. I want to draw the circuit as it is executed on the device by using draw_mpl() with the expasion_strategy='device'. I expect to get a figure of the circuit in terms of CNOT, RZ, X and H gates.

Actual behavior

Instead I get a circuit as if I had not specified the expansion strategy, i.e. a figure one gate on all qubits called UCCSD.

Additional information

When using draw() for the exact same circuit I get the gates that I would have expected.

Source code

from pennylane import qchem
import pennylane as qml
from pennylane import numpy as np

#setup hamiltonian
symbols = ["H", "H"]
coordinates = np.array([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])
h2_ham, num_qubits = qchem.molecular_hamiltonian(symbols, coordinates)

# UCC ansatz for H2 , 3 parameters
electrons = 2
ref_state = qchem.hf_state(electrons, num_qubits) 
singles, doubles = qchem.excitations(electrons, num_qubits)
s_wires, d_wires = qchem.excitations_to_wires(singles, doubles)

def complete_circuit(params, wires):  
    qml.UCCSD(params, wires, s_wires=s_wires,d_wires=d_wires,init_state=ref_state)
    return qml.expval(h2_ham)

devicename="default.qubit"
dev=qml.device(devicename, wires=num_qubits)
wrapped_cost_1 = qml.QNode(complete_circuit, dev)
params = np.random.normal(0, np.pi, len(singles)+len(doubles))

print(qml.draw_mpl(wrapped_cost_1, expansion_strategy='device')(params, range(dev.num_wires)))
print(qml.draw(wrapped_cost_1, expansion_strategy='device')(params, range(dev.num_wires)))

Tracebacks

No response

System information

Name: PennyLane
Version: 0.21.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /home/marita.oliv/Project/venv/lib/python3.8/site-packages
Requires: appdirs, autoray, toml, networkx, semantic-version, retworkx, pennylane-lightning, scipy, numpy, cachetools, autograd
Required-by: PennyLane-qiskit, PennyLane-Qchem, PennyLane-Lightning
Platform info:           Linux-5.4.0-110-generic-x86_64-with-glibc2.29
Python version:          3.8.10
Numpy version:           1.22.2
Scipy version:           1.8.0
Installed devices:
- default.gaussian (PennyLane-0.21.0)
- default.mixed (PennyLane-0.21.0)
- default.qubit (PennyLane-0.21.0)
- default.qubit.autograd (PennyLane-0.21.0)
- default.qubit.jax (PennyLane-0.21.0)
- default.qubit.tf (PennyLane-0.21.0)
- default.qubit.torch (PennyLane-0.21.0)
- qiskit.aer (PennyLane-qiskit-0.21.0)
- qiskit.basicaer (PennyLane-qiskit-0.21.0)
- qiskit.ibmq (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.21.0)
- lightning.qubit (PennyLane-Lightning-0.21.0)
None

Existing GitHub issues

• [X] I have searched existing GitHub issues to make sure the issue does not already exist.

[albi3ro]

Have you tried with a more recent version of PennyLane? The above code works as expected in v0.26.

[qcabepsilon]

It indeed works with a later version of pennylane.