jit with sampling multi-wire observable

[albi3ro]

Context:

Jitting was failing with qml.sample with an observable on more than one wire.

import jax

jax.config.update("jax_enable_x64", True)

dev = qml.device("default.qubit", wires=2, shots=100)

@qml.qnode(dev, interface="jax")
def circuit(x):
    qml.RX(x, wires=0)
    return qml.sample(qml.PauliX(0) @ qml.PauliY(1))

results = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))

Description of the Change:

Update SampleMP.shape so that if there's an observable, there's only one dimension output.

Benefits:

Possible Drawbacks:

Related GitHub Issues:

Fixes #5372 [sc-58779]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.65%. Comparing base (54d43b8) to head (ff81b9f).
Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5422      +/-   ##
==========================================
+ Coverage   99.63%   99.65%   +0.01%     
==========================================
  Files         399      401       +2     
  Lines       37125    36946     -179     
==========================================
- Hits        36990    36817     -173     
+ Misses        135      129       -6     

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

[astralcai]

I'm not sure if I understand the issue. Why would the shape always be 1?

[albi3ro]

I'm not sure if I understand the issue. Why would the shape always be 1?

Because with an observable, we are sampling from the observables eigenvalues. So just one single eigenvalue, even if there's multiple wires.

[astralcai]

Looks good 🐱

[mudit2812]

Looks good! Just one changelog comment.