Repository: PennyLaneAI/pennylane (python)

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.

Pull request description:
Measuring projectors on a state with parameter broadcasting.
This PR supercedes #4971  and fixes #4975 . [sc-52681]

As this is a high priority bug and we have an upcoming release, I am fixing this here.

Nicely this bug can be fixed by removing the specialized logic and just always just the default logic. This will be easier to maintain and extend in the future.

Patch:
diff --git a/doc/releases/changelog-0.34.0.md b/doc/releases/changelog-0.34.0.md
index 21d33279a..2d64b8c32 100644
--- a/doc/releases/changelog-0.34.0.md
+++ b/doc/releases/changelog-0.34.0.md
@@ -708,7 +708,10 @@
   [(#4951)](https://github.com/PennyLaneAI/pennylane/pull/4951)
 
 * `MPLDrawer` does not add the bonus space for classical wires when no classical wires are present.
-  [(#5987)](https://github.com/PennyLaneAI/pennylane/pull/4987)
+  [(#4987)](https://github.com/PennyLaneAI/pennylane/pull/4987)
+
+* `Projector` now works with parameter-broadcasting.
+  [(#4993)](https://github.com/PennyLaneAI/pennylane/pull/4993)
   
 * The jax-jit interface can now be used with float32 mode.
   [(#4990)](https://github.com/PennyLaneAI/pennylane/pull/4990)
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index dcb1957c4..643755dc5 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -52,7 +52,6 @@ from pennylane.measurements import (
     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 @@ class QubitDevice(Device):
         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 @@ class QubitDevice(Device):
         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
index 3f935dc1a..74bf3c5b1 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -19,7 +19,6 @@ from typing import Sequence, Tuple, Union
 
 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 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
         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 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
         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
index 131e153ab..798a47d6c 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -20,7 +20,6 @@ from typing import Sequence, Tuple, Union
 
 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 @@ class VarianceMP(SampleMeasurement, StateMeasurement):
         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 @@ class VarianceMP(SampleMeasurement, StateMeasurement):
         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
index 525567a91..6b5d2a6f9 100644
--- a/tests/ops/qubit/test_observables.py
+++ b/tests/ops/qubit/test_observables.py
@@ -637,6 +637,19 @@ class TestBasisStateProjector:
         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."""