Consider measurements when resolving to jax-jit, fix indexing issues with DeviceArray in QubitDevice

[antalszava]

Context:

import pennylane as qml
import jax

from pennylane import numpy as np

dev = qml.device('default.qubit.jax', wires=2, shots=10)
projector = np.zeros((2 ** 2, 2 ** 2))
projector[0][0] = 1

@jax.jit
@qml.qnode(dev, interface='jax')
def circ(projector):
    qml.PauliZ(wires=0)
    return qml.expval(qml.Hermitian(projector, wires=range(2)))

print(circ(projector))

Raises errors due to:

1. Indexing into the samples generated by QubitDevice and stores under the hood with a DeviceArray instead of a NumPy array. We have a DeviceArray because of using the default.qubit.jax device.
2. Once the 1. issue is solved, issues with concretization arise from JAX due to not using the jax-jit interface. This is an issue in the automatic change to jax-jit when interface="jax" is specified.

Description of the Change:

1. Casts the indices used to a NumPy array;
2. Changes the logic that determines the jax-jit interface to be used such that observables are considered too.

Benefits:

The original snippet executes without issues.

Possible Drawbacks:
N/A

Related GitHub Issues:
Closes #2406

[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

Merging #2427 (26e9388) into master (5c7c77b) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #2427   +/-   ##
=======================================
  Coverage   99.45%   99.45%           
=======================================
  Files         244      244           
  Lines       18976    18978    +2     
=======================================
+ Hits        18872    18874    +2     
  Misses        104      104           

Impacted Files	Coverage Δ
pennylane/_qubit_device.py	98.80% <100.00%> (+<0.01%)	⬆️
pennylane/interfaces/jax.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 5c7c77b...26e9388. Read the comment docs.

[rmoyard]

Thank you @antalszava nice solutions! It looks good to me 💯

[rmoyard]

Very good catch @antalszava 🕵️‍♂️

[antalszava]

[sc-17102]