[BUG] `hamiltonian_expand` produces incorrect dimensions with shot vectors

[astralcai]

Expected behavior

The return type specification states that the output dimension that corresponds to the shot vector should be leading the dimension that corresponds to parameter broadcasting. For example,

import pennylane as qml

dev_with_shot_vector = qml.device("default.qubit", shots=(10, 10, 10))
H = qml.Hamiltonian([1, 2.0], [qml.PauliZ(0), qml.PauliX(0)])

@qml.qnode(dev_with_shot_vector)
def circuit(inputs):
    qml.RX(inputs, wires=0)
    return qml.expval(H)

circuit([0.1, 0.1])

The above circuit produces:

(array([1.4, 1.8]), array([ 1.4, -0.2]), array([1.4, 1. ]))

Actual behavior

When hamiltonian_expand is applied,

@qml.transforms.hamiltonian_expand
@qml.qnode(dev_with_shot_vector)
def circuit(inputs):
    qml.RX(inputs, wires=0)
    return qml.expval(H)

circuit([0.1, 0.1])

The above circuit produces

array([[ 1. ,  1.4,  2.2],
       [-0.2, -0.2,  1. ]])

which is an array of shape (2, 3), where the dimension for shot vectors is trailing, as opposed to leading, the dimension for parameter broadcasting.

Additional information

No response

Source code

No response

Tracebacks

No response

System information

Name: PennyLane
Version: 0.36.0.dev0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/astral.cai/Workspace/pennylane/venv/lib/python3.9/site-packages
Editable project location: /Users/astral.cai/Workspace/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Cirq, PennyLane-qiskit, PennyLane_Lightning
Platform info:           macOS-14.4.1-arm64-arm-64bit
Python version:          3.9.19
Numpy version:           1.26.3
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.36.0.dev0)
- default.gaussian (PennyLane-0.36.0.dev0)
- default.mixed (PennyLane-0.36.0.dev0)
- default.qubit (PennyLane-0.36.0.dev0)
- default.qubit.autograd (PennyLane-0.36.0.dev0)
- default.qubit.jax (PennyLane-0.36.0.dev0)
- default.qubit.legacy (PennyLane-0.36.0.dev0)
- default.qubit.tf (PennyLane-0.36.0.dev0)
- default.qubit.torch (PennyLane-0.36.0.dev0)
- default.qutrit (PennyLane-0.36.0.dev0)
- null.qubit (PennyLane-0.36.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.35.1)
- qiskit.aer (PennyLane-qiskit-0.35.0)
- qiskit.basicaer (PennyLane-qiskit-0.35.0)
- qiskit.ibmq (PennyLane-qiskit-0.35.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.35.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.35.0)
- qiskit.remote (PennyLane-qiskit-0.35.0)
- cirq.mixedsimulator (PennyLane-Cirq-0.35.0.dev0)
- cirq.pasqal (PennyLane-Cirq-0.35.0.dev0)
- cirq.qsim (PennyLane-Cirq-0.35.0.dev0)
- cirq.qsimh (PennyLane-Cirq-0.35.0.dev0)
- cirq.simulator (PennyLane-Cirq-0.35.0.dev0)

Existing GitHub issues

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