[BUG] [Return-types] Update the output shape of qml.probs for default.qubit.torch

[antalszava]

Expected behavior

Using default.qubit.torch and backpropagation, return qml.probs(0), qml.probs(1) returns a (2, 2) shape tensor (just as with TensorFlow).

Actual behavior

A (4,) tensor is returned for return qml.probs(0), qml.probs(1) in Torch+backprop. This could be due to differing behaviour in the DefaultQubitTorch._asarray() method.

Additional information

No response

Source code

import pennylane as qml

dev = qml.device('default.qubit.torch', wires=3)

@qml.qnode(dev, interface='torch')
def circuit():
    return qml.probs(0), qml.probs(1)

circuit()

tensor([1., 0., 1., 0.], dtype=torch.float64)

import pennylane as qml

dev = qml.device('default.qubit.tf', wires=3)

@qml.qnode(dev, interface='tf')
def circuit():
    return qml.probs(0), qml.probs(1)

circuit()

<tf.Tensor: shape=(2, 2), dtype=float64, numpy=
array([[1., 0.],
       [1., 0.]])>

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.18.0.dev0
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: /xanadu/pennylane
Requires: numpy, scipy, networkx, autograd, toml, appdirs, semantic-version, autoray, cachetools
Required-by: PennyLane-Cirq, PennyLane-Orquestra, PennyLane-Qchem, PennyLane-SF, pennylane-qulacs, PennyLane-IonQ, amazon-braket-pennylane-plugin, PennyLane-Forest, PennyLane-Honeywell, PennyLane-qiskit, PennyLane-Lightning
Platform info:           Linux-5.11.0-27-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.20.3
Scipy version:           1.7.1
Installed devices:
- cirq.mixedsimulator (PennyLane-Cirq-0.17.1)
- cirq.pasqal (PennyLane-Cirq-0.17.1)
- cirq.qsim (PennyLane-Cirq-0.17.1)
- cirq.qsimh (PennyLane-Cirq-0.17.1)
- cirq.simulator (PennyLane-Cirq-0.17.1)
- orquestra.forest (PennyLane-Orquestra-0.15.0)
- orquestra.ibmq (PennyLane-Orquestra-0.15.0)
- orquestra.qiskit (PennyLane-Orquestra-0.15.0)
- orquestra.qulacs (PennyLane-Orquestra-0.15.0)
- strawberryfields.fock (PennyLane-SF-0.16.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.16.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.16.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.16.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.16.0.dev0)
- qulacs.simulator (pennylane-qulacs-0.17.0.dev0)
- ionq.qpu (PennyLane-IonQ-0.17.0.dev0)
- ionq.simulator (PennyLane-IonQ-0.17.0.dev0)
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.4.1.dev0)
- forest.numpy_wavefunction (PennyLane-Forest-0.17.0.dev0)
- forest.qvm (PennyLane-Forest-0.17.0.dev0)
- forest.wavefunction (PennyLane-Forest-0.17.0.dev0)
- honeywell.hqs (PennyLane-Honeywell-0.16.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.18.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.18.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.18.0.dev0)
- default.gaussian (PennyLane-0.18.0.dev0)
- default.mixed (PennyLane-0.18.0.dev0)
- default.qubit (PennyLane-0.18.0.dev0)
- default.qubit.autograd (PennyLane-0.18.0.dev0)
- default.qubit.jax (PennyLane-0.18.0.dev0)
- default.qubit.tf (PennyLane-0.18.0.dev0)
- default.qubit.torch (PennyLane-0.18.0.dev0)
- default.tensor (PennyLane-0.18.0.dev0)
- default.tensor.tf (PennyLane-0.18.0.dev0)

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

[dwierichs]

An instance of where this bug became apparent/relevant is the use of broadcasting in qml.gradients.param_shift: Due to the differing output shape of the torch device in its _asarray method, the if-condition in gradients/parameter_shift._evaluate_gradient needs to check for torch devices and handles the corresponding case differently than for all other interfaces. This leads to self-consistent behaviour within the torch interface/device scenario, but showcases that the _asarray method should be adapted to yield consistent results with other devices or produce predictable output shapes instead of iteratively flattening and concatenating the returned values.

[antalszava]

Note that this bug will no longer be relevant with the new return types:

[albi3ro]

Given this bug is fixed by the new return types specification, I'm marking this bug as done.