Allow vector-valued QNodes with JAX using host_callback.call

[antalszava]

Context:

Vector-valued QNodes may include:

1. Multiple scalar return types: e.g., return qml.expval(qml.PauliZ(0), qml.expval(qml.PauliZ(1);
2. QNodes with return types qml.probs, qml.state or qml.density_matrix.

The JAX interface doesn't support these return types. The main reason is that host_callback.call, the underlying function that is being used requires the output shape to be passed. The current logic always considers tapes with scalar outputs.

Uses the machinery introduced in #2044.

Description of the Change:

• Updates the MeasurementProcess class such that shape and numeric_type are methods because there were uncovered edge cases of not having pre-set shapes/numeric types (e.g., self._shape was None);
• Changes the shape that is being passed to host_callback.call.

Benefits:

Vector-valued QNodes can be evaluated using the JAX JIT interface that uses host_callback.call; host_callback.call is jittable.

Possible Drawbacks:

The JAX JIT interface doesn't support jax.jacobian (see discussion in #2163).

Related GitHub Issues:
Closes #1208, #2404

---

Testing:

Testing categories:

• Shape
• Correctness (comparison with the expected result)

Multiple scalar outputs:

• qml.expval, qml.var

Single vector valued outputs:

• qml.probs
• qml.state
• qml.density_matrix

Mixed vector and scalar-valued outputs:

• Scalar valued gradient functions using vector-valued QNodes
• Explicitly consider multi-tape cases
• Check what's wrong with the xfail test case (independent param)
• Double-check if output dimensions are the same as with other interfaces

[codecov]

Codecov Report

Merging #2034 (cb485c3) into master (1b788a6) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #2034   +/-   ##
=======================================
  Coverage   99.61%   99.61%           
=======================================
  Files         251      251           
  Lines       20553    20569   +16     
=======================================
+ Hits        20473    20489   +16     
  Misses         80       80           

Impacted Files	Coverage Δ
pennylane/interfaces/autograd.py	100.00% <ø> (ø)
pennylane/interfaces/jax.py	100.00% <100.00%> (ø)
pennylane/interfaces/jax_jit.py	100.00% <100.00%> (ø)
pennylane/measurements.py	99.62% <100.00%> (+0.01%)	⬆️
pennylane/tape/tape.py	98.89% <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 1b788a6...cb485c3. Read the comment docs.

[maliasadi]

Great job @antalszava! 😍 It's awesome to see the JAX support for vector-value QNodes. I only have a few minor suggestions; but I am happy to approve.

[rmoyard]

Hi @antalszava ! Thanks for this awesome work 💯 No major blockers, but I have some questions. I think you can also update the documentation (chart with configurations) :)

[antalszava]

Thank you @maliasadi, @rmoyard for the reviews! The comments should be addressed now.

(One thing todo for me will be to double-check an edge case of multiple tapes with multiple parameters for forward mode.)

[rmoyard]

Awesome work @antalszava 💯 I've left a comment about the problem you mentioned, we can always come back to it 👍 Thank you for changing VnEntropy and MutualInformation measurements process!