[BUGFIX] shadow_expval doesn't work with parameter broadcasting

[PietropaoloFrisoni]

Context: WIth default.qubit, parameter broadcasting does not work with shadow_expval (or classical_shadow) as the measurement process.

These latest developments might be relevant for historical reasons:

• In this PR, support for parameter broadcasting was added for measure and measure_with_samples (presumably as an effort for the new device API). Specifically, a new keyword argument is_state_batched has been introduced to these functions to deal with the parameter broadcasting case.

• In this PR, the classical shadow measurement process was integrated into the new device API as well. A new logical branch was integrated into measure_with_samples adding a new _measure_classical_shadow function. This function did not receive the is_state_batched argument, so apparently supporting parameter broadcasting for classical shadows was not part of the requirements (the shortcut story was not linked so I cannot investigate further).

• In this PR, the is_state_batched argument has been added to _measure_classical_shadow, but support for parameter broadcasting was never implemented (the argument is unused).

Description of the Change: We added an internal conditional quantum transform to default.qubit such that the qml.transforms.broadcast_expand is conditionally applied if ShadowExpvalMP or ClassicalShadowMP are found in the tape measurement process.

Benefits: Now default.qubit supports (non-native) parameter broadcasting with shadow measurement processes.

Possible Drawbacks: Since we are applying the broadcast_expand quantum transform, we are not supporting native parameter broadcasting, This should not be a big deal.

Related GitHub Issues: #6301

Related ShortCut Stories: [sc-74387]

[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 ✅

Please upload report for BASE (v0.39.0-rc0@92e48d3). Learn more about missing BASE report.

Additional details and impacted files

@@              Coverage Diff               @@
##             v0.39.0-rc0    #6464   +/-   ##
==============================================
  Coverage               ?   99.38%           
==============================================
  Files                  ?      452           
  Lines                  ?    42797           
  Branches               ?        0           
==============================================
  Hits                   ?    42535           
  Misses                 ?      262           
  Partials               ?        0           

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

[mudit2812]

Looks good to me. I don't think there is a way around dispatching pre-processing to the device, but this PR only makes broadcasting compatible with shadow measurements if using default.qubit. We should probably specify that in the changelog.

[lillian542]

LGTM!