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:
github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4993?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> :exclamation: No coverage uploaded for pull request base (`v0.34.0-rc0@db4b762`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).


<details><summary>Additional details and impacted files</summary>


```diff
@@              Coverage Diff               @@
##             v0.34.0-rc0    #4993   +/-   ##
==============================================
  Coverage               ?   99.66%           
==============================================
  Files                  ?      392           
  Lines                  ?    35260           
  Branches               ?        0           
==============================================
  Hits                   ?    35142           
  Misses                 ?      118           
  Partials               ?        0           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4993?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

albi3ro: > how confident are you that this won't break plugins?

I think its much more likely that this will fix any plugins that are also having this issue.

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."""