Repository: PennyLaneAI/pennylane

Issue:
[BUG] qubit device expval method doesn't work with batch and BasisStateProjector
### Expected behavior

When utilizing a qubit device with batch processing, expval output, and BasisStateProjector, the probabilities need to be correctly indexed in the good dimension.

### Actual behavior

When utilizing a qubit device with batch processing, expval output, and BasisStateProjector, the probabilities are incorrectly index in the wrong dimension. This result in the expval output returning the probabilities of a vector within the batch, rather than providing the probabilities of the desired state for all entries in the batch.

### Additional information

A PR is already open for this issue. See https://github.com/PennyLaneAI/pennylane/pull/4971 .

### Source code

```shell
import pennylane as qml
from pennylane.devices import DefaultQubit
from pennylane.ops.qubit.observables import BasisStateProjector
import numpy as np


def circuit(x, y):
    n_wires = qml.math.shape(x)[-1]
    for i in range(n_wires):
        qml.RX(x[:, i], wires=i)
        qml.RY(y[:, i], wires=i)
    projector: BasisStateProjector = qml.Projector(np.zeros(n_wires), wires=range(n_wires))
    return qml.expval(projector)


if __name__ == '__main__':
    device = DefaultQubit(wires=2)
    qnode = qml.QNode(circuit, device)
    batch_size = device.num_wires * 3 + 1
    x = np.random.uniform(0, 2 * np.pi, size=(batch_size, device.num_wires))
    y = np.random.uniform(0, 2 * np.pi, size=(batch_size, device.num_wires))
    out = qnode(x, y)
    np.testing.assert_allclose(qml.math.shape(out), (batch_size,))
```


### Tracebacks

```shell
Traceback (most recent call last):
  File "E:\Github\FermionicSimulation\scratches\pennylane_pr.py", line 23, in <module>
    np.testing.assert_allclose(qml.math.shape(out), (batch_size,))
  File "E:\Github\FermionicSimulation\venv\lib\site-packages\numpy\testing\_private\utils.py", line 1527, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "E:\Github\FermionicSimulation\venv\lib\site-packages\numpy\testing\_private\utils.py", line 844, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=1e-07, atol=0

Mismatched elements: 1 / 1 (100%)
Max absolute difference: 3
Max relative difference: 0.42857143
 x: array([4])
 y: array([7])
```


### System information

```shell
Name: PennyLane
Version: 0.32.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: e:\github\fermionicsimulation\venv\lib\site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info:           Windows-10-10.0.22621-SP0
Python version:          3.9.12
Numpy version:           1.23.5
Scipy version:           1.11.3
Installed devices:
- default.gaussian (PennyLane-0.32.0)
- default.mixed (PennyLane-0.32.0)
- default.qubit (PennyLane-0.32.0)
- default.qubit.autograd (PennyLane-0.32.0)
- default.qubit.jax (PennyLane-0.32.0)
- default.qubit.tf (PennyLane-0.32.0)
- default.qubit.torch (PennyLane-0.32.0)
- default.qutrit (PennyLane-0.32.0)
- null.qubit (PennyLane-0.32.0)
- lightning.qubit (PennyLane-Lightning-0.32.0)
```


### Existing GitHub issues

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


Discussion:

Resolving patch:
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -52,7 +52,6 @@
     VnEntropyMP,
     Shots,
 )
-from pennylane.ops.qubit.observables import BasisStateProjector
 from pennylane.resource import Resources
 from pennylane.operation import operation_derivative, Operation
 from pennylane.tape import QuantumTape
@@ -1307,14 +1306,6 @@ def marginal_prob(self, prob, wires=None):
         return self._reshape(prob, flat_shape)
 
     def expval(self, observable, shot_range=None, bin_size=None):
-        if isinstance(observable, BasisStateProjector):
-            # branch specifically to handle the basis state projector observable
-            idx = int("".join(str(i) for i in observable.parameters[0]), 2)
-            probs = self.probability(
-                wires=observable.wires, shot_range=shot_range, bin_size=bin_size
-            )
-            return probs[idx]
-
         # exact expectation value
         if self.shots is None:
             try:
@@ -1343,14 +1334,6 @@ def expval(self, observable, shot_range=None, bin_size=None):
         return np.squeeze(np.mean(samples, axis=axis))
 
     def var(self, observable, shot_range=None, bin_size=None):
-        if isinstance(observable, BasisStateProjector):
-            # branch specifically to handle the basis state projector observable
-            idx = int("".join(str(i) for i in observable.parameters[0]), 2)
-            probs = self.probability(
-                wires=observable.wires, shot_range=shot_range, bin_size=bin_size
-            )
-            return probs[idx] - probs[idx] ** 2
-
         # exact variance value
         if self.shots is None:
             try:
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -19,7 +19,6 @@
 
 import pennylane as qml
 from pennylane.operation import Operator
-from pennylane.ops.qubit.observables import BasisStateProjector
 from pennylane.wires import Wires
 
 from .measurements import Expectation, SampleMeasurement, StateMeasurement
@@ -107,15 +106,6 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        if isinstance(self.obs, BasisStateProjector):
-            # branch specifically to handle the basis state projector observable
-            idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
-            with qml.queuing.QueuingManager.stop_recording():
-                probs = qml.probs(wires=self.wires).process_samples(
-                    samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
-                )
-            return probs[idx]
-
         # estimate the ev
         op = self.mv if self.mv is not None else self.obs
         with qml.queuing.QueuingManager.stop_recording():
@@ -130,15 +120,6 @@ def process_samples(
         return qml.math.squeeze(qml.math.mean(samples, axis=axis))
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
-        if isinstance(self.obs, BasisStateProjector):
-            # branch specifically to handle the basis state projector observable
-            idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
-            with qml.queuing.QueuingManager.stop_recording():
-                probs = qml.probs(wires=self.wires).process_state(
-                    state=state, wire_order=wire_order
-                )
-            return probs[idx]
-
         # This also covers statistics for mid-circuit measurements manipulated using
         # arithmetic operators
         eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -20,7 +20,6 @@
 
 import pennylane as qml
 from pennylane.operation import Operator
-from pennylane.ops.qubit.observables import BasisStateProjector
 from pennylane.wires import Wires
 
 from .measurements import SampleMeasurement, StateMeasurement, Variance
@@ -107,16 +106,6 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        if isinstance(self.obs, BasisStateProjector):
-            # branch specifically to handle the basis state projector observable
-            idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
-            # we use ``self.wires`` instead of ``self.obs`` because the observable was
-            # already applied before the sampling
-            probs = qml.probs(wires=self.wires).process_samples(
-                samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
-            )
-            return probs[idx] - probs[idx] ** 2
-
         # estimate the variance
         op = self.mv if self.mv is not None else self.obs
         with qml.queuing.QueuingManager.stop_recording():
@@ -131,17 +120,6 @@ def process_samples(
         return qml.math.squeeze(qml.math.var(samples, axis=axis))
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
-        if isinstance(self.obs, BasisStateProjector):
-            # branch specifically to handle the basis state projector observable
-            idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
-            # we use ``self.wires`` instead of ``self.obs`` because the observable was
-            # already applied to the state
-            with qml.queuing.QueuingManager.stop_recording():
-                probs = qml.probs(wires=self.wires).process_state(
-                    state=state, wire_order=wire_order
-                )
-            return probs[idx] - probs[idx] ** 2
-
         # This also covers statistics for mid-circuit measurements manipulated using
         # arithmetic operators
         eigvals = qml.math.asarray(self.eigvals(), dtype="float64")


Test patch:
diff --git a/tests/ops/qubit/test_observables.py b/tests/ops/qubit/test_observables.py
--- a/tests/ops/qubit/test_observables.py
+++ b/tests/ops/qubit/test_observables.py
@@ -637,6 +637,19 @@ def test_matrix_representation(self, basis_state, expected, n_wires, tol):
         assert np.allclose(res_dynamic, expected, atol=tol)
         assert np.allclose(res_static, expected, atol=tol)
 
+    @pytest.mark.parametrize(
+        "dev", (qml.device("default.qubit"), qml.device("default.qubit.legacy", wires=1))
+    )
+    def test_integration_batched_state(self, dev):
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.Projector([0], wires=0))
+
+        x = np.array([0.4, 0.8, 1.2])
+        res = circuit(x)
+        assert qml.math.allclose(res, np.cos(x / 2) ** 2)
+
 
 class TestStateVectorProjector:
     """Tests for state vector projector observable."""