dynamic_one_shot supports broadcasting; broadcast_expand supports shot vectors

[mudit2812]

Context:
Native mid-circuit measurements with default.qubit are not compatible with parameter broadcasting. Due to the complexity of a "native" implementation, I decided to use broadcast_expand, but realized that it does not work with shot vectors.

Description of the Change:
This PR does two things:

• Update dynamic_one_shot transform to use broadcast_expand and process batched results correctly.
• Update broadcast_expand to support shot vectors.
• Raise error when postselecting with broadcasting and returning samples. This change was made to both dynamic_one_shot and defer_measurements because both transforms use broadcast_expand for broadcasting, although defer_measurements only uses broadcast_expand with postselection.

Note that broadcasting with qml.sample and postselection will still not work due to ragged dimensions. If reviewers are okay with it, I would like to merge this and leave that improvement as technical debt. cc @trbromley @isaacdevlugt.
Edit about note: Talked offline, decided to raise a more informative error if a user requests postselection with broadcasting when returning samples.

Benefits:
Both transforms are more capable.

Possible Drawbacks:
Because of the stacking and squeezing happening in the post-processing function of broadcast_expand, counts dictionaries get wrapped inside 0-D numpy arrays, which makes indexing into the dict impossible. To access the dictionary and its contents, one has to use array.item() to extract the single item inside the array.

Related GitHub Issues:
#5443

[mudit2812]

[sc-60581] [sc-59941]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.66%. Comparing base (a47d9bc) to head (7569291).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5473      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         406      406              
  Lines       37881    37623     -258     
==========================================
- Hits        37758    37498     -260     
- Misses        123      125       +2     

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

[dwierichs]

Just some small comments/requests, and a question about those wrapped dicts.
Very nice extension of the recursion logic both in code and tests 💯 💪
And great testing @mudit2812 🎉

[vincentmr]

validate_measurements needs a small fix, but otherwise looks good.

[vincentmr]

LGMT, good job @mudit2812 !

[dwierichs]

Thanks for the additional fix of Conditional @mudit2812 ! 💯

Is there a test that would fail without that fix?

Maybe you could also mention the fix of Conditional in the changelog? And shouldn't there be some unit tests of Conditional affected or bind_new_parameters affected somehow? 🙃

[dwierichs]

LGTM!