fix jax jit with broadcasted probs

[albi3ro]

Fixes Issue #4738

The ProbabilityMP.shape did not take into account probabilities that broadcasted across all available wires.

[sc-48626]

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit 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]

Codecov Report

All modified and coverable lines are covered by tests ✅

❗ No coverage uploaded for pull request base (v0.33.0-rc0@cd72f35). Click here to learn what that means.

Additional details and impacted files

@@              Coverage Diff               @@
##             v0.33.0-rc0    #4742   +/-   ##
==============================================
  Coverage               ?   99.64%           
==============================================
  Files                  ?      380           
  Lines                  ?    33995           
  Branches               ?        0           
==============================================
  Hits                   ?    33873           
  Misses                 ?      122           
  Partials               ?        0           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mudit2812]

This problem exists for sample().shape and state().shape as well. I think it's worthwhile updating them in this PR as well.

[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. I don't quite want to make a change like that this week.

[mudit2812]

Thanks!

[timmysilv]

@albi3ro we might not need a breaking change. I used to just replace device with some object that also had wires defined (at the time, it was self) but we can do anything. I added it early in the device port, but I removed it from qscript.py here