Improving support of counts and sample with default.mixed

[obliviateandsurrender]

Context: Counts should work with interfaces (torch, tf, jax) and default.mixed but it is not the case. This PR fixes it.

Description of the Change:

Benefits:

Possible Drawbacks:

Related GitHub Issues: #2984

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

[obliviateandsurrender]

[sc-25027]

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.66%. Comparing base (016d2cc) to head (0f47690).
Report is 3 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5514      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         406      406              
  Lines       37928    37643     -285     
==========================================
- Hits        37804    37518     -286     
- Misses        124      125       +1     

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

[mudit2812]

LGTM!

[Jaybsoni]

Looks good 👍🏼,

This might also natively support jax now. If that's the case it would be nice to also add jax tests. Not blocking though.

Edit:

I realized afterwards that this fix is for the torch default mixed device. Looks good 👍🏼