Repository: PennyLaneAI/pennylane
Pull request #4742 (2023-10-26 14:57:11)

fix jax jit with broadcasted probs
Fixes Issue #4738 

The `ProbabilityMP.shape` did not take into account probabilities that broadcasted across all available wires.

[sc-48626]

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/4742?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.33.0-rc0@a9ad7a0`). [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.33.0-rc0    #4742   +/-   ##
==============================================
  Coverage               ?   99.64%           
==============================================
  Files                  ?      380           
  Lines                  ?    33982           
  Branches               ?        0           
==============================================
  Hits                   ?    33860           
  Misses                 ?      122           
  Partials               ?        0           
```




</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4742?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: > This problem exists for `sample().shape` and `state().shape` as well. I think it's worthwhile updating them in this PR as well.

jax-jit does not work with `sample` or `state` anyway, but it looks like they don't have the problem:
```
>>> qml.state().shape(qml.device('default.qubit', wires=(0,1)), qml.measurements.Shots(None))
(4,)
>>> qml.sample().shape(qml.device('default.qubit', wires=(0,1)), qml.measurements.Shots(5))
(5, 2)
```
mudit2812: > jax-jit does not work with sample or state anyway, but it looks like they don't have the problem:

Sorry, I meant for cases when devices without wires are used.
albi3ro: > > jax-jit does not work with sample or state anyway, but it looks like they don't have the problem:
> 
> Sorry, I meant for cases when devices without wires are used.

For devices without wires we are going to need to make a breaking change to the call signature of `MeasurementProcess.shape`. We will need to change from `(device, shots)` to `(wires, shots)` as the wires may be coming from the tape instead of from the device.

Patch:
diff --git a/doc/releases/changelog-0.33.0.md b/doc/releases/changelog-0.33.0.md
index 43194290b..1ea85bd82 100644
--- a/doc/releases/changelog-0.33.0.md
+++ b/doc/releases/changelog-0.33.0.md
@@ -633,6 +633,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Jax jit now works when a probability measurement is broadcasted onto all wires.
+  [(#4742)](https://github.com/PennyLaneAI/pennylane/pull/4742)
+
 * Fixed `LocalHilbertSchmidt.compute_decomposition` so that the template can be used in a QNode.
   [(#4719)](https://github.com/PennyLaneAI/pennylane/pull/4719)
 
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 9b5cbf9a4..ae1dc5315 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -148,6 +148,8 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
             sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
         )
         len_wires = len(self.wires)
+        if len_wires == 0:
+            len_wires = len(device.wires) if device.wires else 0
         dim = self._get_num_basis_states(len_wires, device)
 
         return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements))


Test patch:
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
index 2c48a1b97..c209d01fb 100644
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -679,6 +679,7 @@ class TestVectorValuedJIT:
         [
             ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (), tuple),
             ([qml.probs(wires=[0, 1])], (4,), jax.numpy.ndarray),
+            ([qml.probs()], (4,), jax.numpy.ndarray),
         ],
     )
     def test_shapes(self, execute_kwargs, ret_type, shape, expected_type):
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 3cf2c177c..94090002c 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -73,6 +73,16 @@ class TestProbs:
         res = qml.probs(wires=wires)
         assert res.shape(dev, Shots(shots)) == (2 ** len(wires),)
 
+    def test_shape_empty_wires(self):
+        """Test that shape works when probs is broadcasted onto all available wires."""
+        dev = qml.device("default.qubit", wires=(1, 2, 3))
+        res = qml.probs()
+        assert res.shape(dev, Shots(None)) == (8,)
+
+        dev2 = qml.device("default.qubit")
+        res = qml.probs()
+        assert res.shape(dev2, Shots(None)) == (1,)
+
     @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]])
     def test_shape_shot_vector(self, wires):
         """Test that the shape is correct with the shot vector too."""