[BUG] No type safety when passing argument to @qml.device() decorated function

Expected behavior

It would be really nice if I could combine circuits in an obvious way, but the lack of type safety prevents me from doing so whenever the function is decorated with @qml.qnode(device). I really simple example is preparing a VQE circuit from a UCCSD ansatz plus grouped measurement.

I understand that using the built-in methods should take care of most of this for me, but for the purposes of what I am trying to do that simply won’t do.

As a simple example, consider H2 with qubit-wise commuting groups.

import pennylane as qml
from pennylane import qchem
import numpy as np
from functools import partial
from pennylane.templates import UCCSD

#MOLECULE
symbols = ["H", "H"]
coordinates = np.array([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])
h2_ham, n_qubits = qchem.molecular_hamiltonian(symbols, coordinates)
n_electrons = 2

groupings, coeffs = qml.grouping.group_observables(h2_ham.terms[1], h2_ham.terms[0], grouping_type = 'qwc')
dev_noisy = qml.device('default.mixed', wires = n_qubits, shots = 1000)

#prepare UCCSD ansatz
singles, doubles = qchem.excitations(n_electrons, n_qubits)
s_wires, d_wires = qchem.excitations_to_wires(singles, doubles)
ref_state = qchem.hf_state(n_electrons, n_qubits)
ansatz = partial(UCCSD, init_state = ref_state, s_wires = s_wires, d_wires = d_wires)

@qml.qnode(dev_noisy)
def VQE_circuit(params, group, n_qubits):
    print(type(group))
    qml.apply(ansatz(params, wires = range(n_qubits)))
    #FAILS HERE, tensor has no attribute wire
    rotations = qml.grouping.diagonalize_qwc_pauli_words(group)[0]
    for rotation in rotations:
        qml.apply(rotation)
    return qml.sample(wires = range(n_qubits))
params = np.random.randn(len(s_wires) + len(d_wires))

circuits = [VQE_circuit(params, group, n_qubits) for group in groupings]

Actual behavior

However, the lack of type safety prevents me from doing so.

Prior to passing group to VQE_circuit, type(group) = <class 'list'> with a particular object being something like: [PauliY(wires=[0]) @ PauliY(wires=[1]) @ PauliX(wires=[2]) @ PauliX(wires=[3])].

However, once it’s in the body of the function: type(group) = <class 'pennylane.numpy.tensor.tensor'> which poses a problem since that class doesn’t interface with diagonalize_qwc_pauli_words bc pennylane.numpy.tensor.tensor objects don’t have wires.

Note that this all works perfectly fine once the decorator @qml.device is removed.

Additional information

No response

Source code

No response

Tracebacks

No response

System information

WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.
Name: PennyLane
Version: 0.17.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/kharazi/.local/lib/python3.9/site-packages
Requires: networkx, toml, semantic-version, appdirs, autograd, scipy, autoray, numpy
Required-by: VQELib, PennyLane-Qchem
Platform info:           Linux-5.11.0-34-generic-x86_64-with-glibc2.33
Python version:          3.9.5
Numpy version:           1.21.2
Scipy version:           1.7.1
Installed devices:
- default.gaussian (PennyLane-0.17.0)
- default.mixed (PennyLane-0.17.0)
- default.qubit (PennyLane-0.17.0)
- default.qubit.autograd (PennyLane-0.17.0)
- default.qubit.jax (PennyLane-0.17.0)
- default.qubit.tf (PennyLane-0.17.0)
- default.tensor (PennyLane-0.17.0)
- default.tensor.tf (PennyLane-0.17.0)

• I have searched exisisting GitHub issues to make sure the issue does not already exist.

[josh146]

No worries @kharazity! Since this is a duplicate of an existing bug, I will close this in favour of #1459. However, don’t hesitate to open a new issue if you discover any other bugs!

In addition, we also have a PennyLane discussion forum, where you can ask general code and QML related questions: https://discuss.pennylane.ai

[josh146]

@kharazity I’ve opened a PR, #1680, to fix this latter bug 😃