Repository: PennyLaneAI/pennylane
Pull request #2034 (2021-12-15 22:41:07)

Allow vector-valued QNodes with JAX using `host_callback.call`
**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 https://github.com/PennyLaneAI/pennylane/issues/2163).

**Related GitHub Issues:**
Closes https://github.com/PennyLaneAI/pennylane/issues/1208, https://github.com/PennyLaneAI/pennylane/issues/2404

----

Testing:

Testing categories:

* Shape
* Correctness (comparison with the expected result)

Multiple scalar outputs:

- [X] qml.expval, qml.var

Single vector valued outputs:
- [X] qml.probs
- [X] qml.state
- [X] qml.density_matrix

Mixed vector and scalar-valued outputs:

- [x] Scalar valued gradient functions using vector-valued QNodes
- [x] Explicitly consider multi-tape cases
- [x] Check what's wrong with the xfail test case (independent param)
- [x] Double-check if output dimensions are the same as with other interfaces

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/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[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2034](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c09e946) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1389303570e1e11079cf68f755412b00b4095674?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1389303) will **decrease** coverage by `2.21%`.
> The diff coverage is `0.00%`.

> :exclamation: Current head c09e946 differs from pull request most recent head 2902a73. Consider uploading reports for the commit 2902a73 to get more accurate results
[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2034      +/-   ##
==========================================
- Coverage   99.15%   96.93%   -2.22%     
==========================================
  Files         221      221              
  Lines       16844    16859      +15     
==========================================
- Hits        16702    16343     -359     
- Misses        142      516     +374     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.21%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `80.00% <0.00%> (-20.00%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `81.81% <0.00%> (-18.19%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `86.84% <0.00%> (-13.16%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| [pennylane/collections/sum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2NvbGxlY3Rpb25zL3N1bS5weQ==) | `89.47% <0.00%> (-10.53%)` | :arrow_down: |
| ... and [13 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [1389303...2902a73](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2034?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: Hey @antalszava! Has this PR been superseded?
antalszava: Hi Josh, no, this would be the step after having #2044 in. Once we have `QuantumTape.get_output_shape` and `QuantumTape.get_output_domain`, defining the shape to be passed to `call` will be simple and can be more general.
antalszava: [sc-17241]
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.)

Patch:
diff --git a/doc/introduction/interfaces/jax.rst b/doc/introduction/interfaces/jax.rst
index e28624236..dcbb2e46d 100644
--- a/doc/introduction/interfaces/jax.rst
+++ b/doc/introduction/interfaces/jax.rst
@@ -137,14 +137,17 @@ the ``@jax.jit`` decorator can be directly applied to the QNode.
     the computation was just-in-time compiled. This is done by checking if any
     of the input parameters were subject to a JAX transformation. If so, a
     variant of the interface that supports the just-in-time compilation of
-    scalar-valued QNodes (i.e., those that have a single expectation value or
-    variance measurement) will be used. This is equivalent to passing
-    ``interface='jax-jit'``.
+    QNodes will be used. This is equivalent to passing ``interface='jax-jit'``.
+
+    Computing the jacobian of vector-valued QNodes is not supported with the
+    JAX JIT interface. The output of vector-valued QNodes can, however, be used
+    in the definition of scalar-valued cost functions whose gradients can be
+    computed.
 
     Specify ``interface='jax-python'`` to enforce support for computing the
-    forward and backward pass of vector-valued QNodes (e.g., QNodes with
-    probability, state or multiple expectation value measurements). This
-    option does not support just-in-time compilation.
+    backward pass of vector-valued QNodes (e.g., QNodes with probability, state
+    or multiple expectation value measurements). This option does not support
+    just-in-time compilation.
 
 
 Randomness: Shots and Samples
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 90f89843c..622f2ed50 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,52 @@
 
 <h3>New features since last release</h3>
 
+* The JAX JIT interface now supports evaluating vector-valued QNodes
+  enabling new types of workflows to utilize the power of just-in-time
+  compilation for significant performance boosts.
+  [(#2034)](https://github.com/PennyLaneAI/pennylane/pull/2034)
+
+  Vector-valued QNodes include those with:
+  * `qml.probs`;
+  * `qml.state`;
+  * `qml.sample` or
+  * multiple `qml.expval` / `qml.var` measurements.
+
+  Consider a QNode that returns basis-state probabilities:
+  ```python
+  dev = qml.device('default.qubit', wires=2)
+  x = jnp.array(0.543)
+  y = jnp.array(-0.654)
+
+  @jax.jit
+  @qml.qnode(dev, diff_method="parameter-shift", interface="jax")
+  def circuit(x, y):
+      qml.RX(x, wires=[0])
+      qml.RY(y, wires=[1])
+      qml.CNOT(wires=[0, 1])
+      return qml.probs(wires=[1])
+  ```
+  The QNode can now be evaluated:
+  ```pycon
+  >>> circuit(x, y)
+  DeviceArray([0.8397495 , 0.16025047], dtype=float32)
+  ```
+  Computing the jacobian of vector-valued QNodes is not supported with the JAX
+  JIT interface. The output of vector-valued QNodes can, however, be used in
+  the definition of scalar-valued cost functions whose gradients can be
+  computed.
+
+  For example, one can define a cost function that outputs the first element of
+  the probability vector:
+  ```python
+  def cost(x, y):
+      return circuit(x, y)[0]
+  ```
+  ```pycon
+  >>> jax.grad(cost, argnums=[0])(x, y)
+  (DeviceArray(-0.2050439, dtype=float32),)
+  ```
+
 * A new quantum information module is added. It includes a function for computing the reduced density matrix functions
   for state vectors and density matrices.
 
@@ -663,6 +709,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Amintor Dusko, Ankit Khandelwal, Avani Bhardwaj, Chae-Yeun Park, Christian Gogolin, Christina Lee, David Wierichs, Edward Jiang, Guillermo Alonso-Linaje,
-Jay Soni, Juan Miguel Arrazola, Katharine Hyatt, Korbinian Kottmann, Maria Schuld, Mikhail Andrenkov, Romain Moyard,
-Qi Hu, Samuel Banning, Soran Jahangiri, Utkarsh Azad, WingCode
+Amintor Dusko, Ankit Khandelwal, Avani Bhardwaj, Chae-Yeun Park, Christian Gogolin, Christina Lee, David Wierichs,
+Edward Jiang, Guillermo Alonso-Linaje, Jay Soni, Juan Miguel Arrazola, Katharine Hyatt, Korbinian Kottmann,
+Maria Schuld, Mikhail Andrenkov, Romain Moyard, Qi Hu, Samuel Banning, Soran Jahangiri, Utkarsh Azad, Antal Száva,
+WingCode
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index e4822635e..97d322918 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -58,7 +58,7 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
         tape.trainable_params = qml.math.get_trainable_indices(params)
 
     # pylint misidentifies autograd.builtins as a dict
-    # pylint:disable=no-member
+    # pylint: disable=no-member
     parameters = autograd.builtins.tuple(
         [autograd.builtins.list(t.get_parameters()) for t in tapes]
     )
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index cd1b58fe5..0917e04ba 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -260,6 +260,36 @@ def _execute(
     return wrapped_exec(params)
 
 
+def _raise_vector_valued_fwd(tapes):
+    """Raises an error for vector-valued tapes in forward mode due to incorrect
+    results being produced.
+
+    There is an issue when jax.jacobian is being used, either due to issues
+    with tensor updating (TypeError: Updates tensor must be of rank 0; got 1)
+    or because jax.vmap introduces a redundant dimensionality in the result by
+    duplicating entries.
+
+    Example to the latter:
+
+    1. Output when using jax.jacobian:
+    DeviceArray([[-0.09983342,  0.01983384],\n
+                 [-0.09983342, 0.01983384]], dtype=float64),
+    DeviceArray([[ 0.        , -0.97517033],\n
+                 [ 0.        , -0.97517033]], dtype=float64)),
+
+    2. Expected output:
+    DeviceArray([[-0.09983342, 0.01983384],\n
+                [ 0.        , -0.97517033]]
+
+    The output produced by this function matches 1.
+    """
+    scalar_outputs = all(t.output_dim == 1 for t in tapes)
+    if not scalar_outputs:
+        raise InterfaceUnsupportedError(
+            "Computing the jacobian of vector-valued tapes is not supported currently in forward mode."
+        )
+
+
 def _execute_with_fwd(
     params,
     tapes=None,
@@ -298,28 +328,7 @@ def _execute_with_fwd(
         # Use the jacobian that was computed on the forward pass
         jacs, params = params
 
-        # Note: there is an issue when jax.jacobian is being used, either due
-        # to issues with tensor updating (TypeError: Updates tensor must be of
-        # rank 0; got 1) or because jax.vmap introduces a redundant
-        # dimensionality in the result by duplicating entries
-        # Example to the latter:
-        #
-        # 1. Output when using jax.jacobian:
-        # DeviceArray([[-0.09983342,  0.01983384],\n
-        #              [-0.09983342, 0.01983384]], dtype=float64),
-        # DeviceArray([[ 0.        , -0.97517033],\n
-        #              [ 0.        , -0.97517033]], dtype=float64)),
-        #
-        # 2. Expected output:
-        # DeviceArray([[-0.09983342, 0.01983384],\n
-        #             [ 0.        , -0.97517033]]
-        #
-        # The output produced by this function matches 2.
-        scalar_outputs = all(t.output_dim == 1 for t in tapes)
-        if not scalar_outputs:
-            raise InterfaceUnsupportedError(
-                "Computing the jacobian of vector-valued tapes is not supported currently in forward mode."
-            )
+        _raise_vector_valued_fwd(tapes)
 
         # Adjust the structure of how the jacobian is returned to match the
         # non-forward mode cases
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
index 83b694a9d..d40d26495 100644
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -23,9 +23,8 @@ from jax.experimental import host_callback
 
 import numpy as np
 import pennylane as qml
-
-from pennylane.measurements import Variance, Expectation, VnEntropy, MutualInfo
 from pennylane.interfaces import InterfaceUnsupportedError
+from pennylane.interfaces.jax import _raise_vector_valued_fwd
 
 dtype = jnp.float64
 
@@ -62,8 +61,6 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
     if max_diff > 1:
         raise InterfaceUnsupportedError("The JAX interface only supports first order derivatives.")
 
-    _validate_tapes(tapes)
-
     for tape in tapes:
         # set the trainable parameters
         params = tape.get_parameters(trainable_only=False)
@@ -92,31 +89,39 @@ def execute(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_d
     )
 
 
-def _validate_tapes(tapes):
-    """Validates that the input tapes are compatible with JAX support.
+def _numeric_type_to_dtype(numeric_type):
+    """Auxiliary function for converting from Python numeric types to JAX
+    dtypes based on the precision defined for the interface."""
+
+    single_precision = dtype is jnp.float32
+    if numeric_type is int:
+        return jnp.int32 if single_precision else jnp.int64
+
+    if numeric_type is float:
+        return jnp.float32 if single_precision else jnp.float64
+
+    # numeric_type is complex
+    return jnp.complex64 if single_precision else jnp.complex128
 
-    Raises:
-        ValueError: if tapes with non-scalar outputs were provided or a return
-        type other than variance and expectation value was used
+
+def _extract_shape_dtype_structs(tapes, device):
+    """Auxiliary function for defining the jax.ShapeDtypeStruct objects given
+    the tapes and the device.
+
+    The host_callback.call function expects jax.ShapeDtypeStruct objects to
+    describe the output of the function call.
     """
+    shape_dtypes = []
+
     for t in tapes:
+        shape = t.shape(device)
 
-        if len(t.observables) != 1:
-            raise InterfaceUnsupportedError(
-                "The jittable JAX interface currently only supports quantum nodes with a single return type."
-            )
+        tape_dtype = _numeric_type_to_dtype(t.numeric_type)
+        shape_and_dtype = jax.ShapeDtypeStruct(tuple(shape), tape_dtype)
 
-        for o in t.observables:
-            return_type = o.return_type
-            if (
-                return_type is not Variance
-                and return_type is not Expectation
-                and return_type is not VnEntropy
-                and return_type is not MutualInfo
-            ):
-                raise InterfaceUnsupportedError(
-                    f"Only Variance and Expectation returns are supported for the jittable JAX interface, given {return_type}."
-                )
+        shape_dtypes.append(shape_and_dtype)
+
+    return shape_dtypes
 
 
 def _execute(
@@ -128,8 +133,6 @@ def _execute(
     gradient_kwargs=None,
     _n=1,
 ):  # pylint: disable=dangerous-default-value,unused-argument
-    # Only have scalar outputs
-    total_size = len(tapes)
     total_params = np.sum([len(p) for p in params])
 
     # Copy a given tape with operations and set parameters
@@ -147,8 +150,8 @@ def _execute(
                 res, _ = execute_fn(new_tapes, **gradient_kwargs)
             return res
 
-        shapes = [jax.ShapeDtypeStruct((1,), dtype) for _ in range(total_size)]
-        res = host_callback.call(wrapper, params, result_shape=shapes)
+        shape_dtype_structs = _extract_shape_dtype_structs(tapes, device)
+        res = host_callback.call(wrapper, params, result_shape=shape_dtype_structs)
         return res
 
     def wrapped_exec_fwd(params):
@@ -213,7 +216,10 @@ def _execute(
                 jacs = gradient_fn(new_tapes, **gradient_kwargs)
             return jacs
 
-        shapes = [jax.ShapeDtypeStruct((1, len(p)), dtype) for p in params]
+        shapes = [
+            jax.ShapeDtypeStruct((len(t.measurements), len(p)), dtype)
+            for t, p in zip(tapes, params)
+        ]
         jacs = host_callback.call(jacs_wrapper, params, result_shape=shapes)
         vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(g, jacs)]
         res = [[jnp.array(p) for p in v] for v in vjps]
@@ -232,10 +238,6 @@ def _execute_with_fwd(
     gradient_kwargs=None,
     _n=1,
 ):  # pylint: disable=dangerous-default-value,unused-argument
-
-    # Only have scalar outputs
-    total_size = len(tapes)
-
     @jax.custom_vjp
     def wrapped_exec(params):
         def wrapper(p):
@@ -251,12 +253,24 @@ def _execute_with_fwd(
             # On the forward execution return the jacobian too
             return res, jacs
 
-        fwd_shapes = [jax.ShapeDtypeStruct((1,), dtype) for _ in range(total_size)]
-        jacobian_shape = [jax.ShapeDtypeStruct((1, len(p)), dtype) for p in params]
+        fwd_shape_dtype_struct = _extract_shape_dtype_structs(tapes, device)
+
+        jacobian_shape = [t.shape(device) + (len(p),) for t in tapes for p in params]
+        jac_dtypes = [_numeric_type_to_dtype(t.numeric_type) for t in tapes]
+
+        # Note: for qml.probs we'll first have a [1,dim] shape for the tape
+        # which is then reduced by the QNode
+        jacobian_shape = [
+            jax.ShapeDtypeStruct(tuple([shape]), dtype)
+            if isinstance(shape, int)
+            else jax.ShapeDtypeStruct(tuple(shape), dtype)
+            for shape, dtype in zip(jacobian_shape, jac_dtypes)
+        ]
+
         res, jacs = host_callback.call(
             wrapper,
             params,
-            result_shape=tuple([fwd_shapes, jacobian_shape]),
+            result_shape=tuple([fwd_shape_dtype_struct, jacobian_shape]),
         )
         return res, jacs
 
@@ -266,6 +280,8 @@ def _execute_with_fwd(
 
     def wrapped_exec_bwd(params, g):
 
+        _raise_vector_valued_fwd(tapes)
+
         # Use the jacobian that was computed on the forward pass
         jacs, params = params
 
diff --git a/pennylane/measurements.py b/pennylane/measurements.py
index 34e43cee3..c347829dc 100644
--- a/pennylane/measurements.py
+++ b/pennylane/measurements.py
@@ -99,16 +99,6 @@ class MeasurementProcess:
             This can only be specified if an observable was not provided.
         eigvals (array): A flat array representing the eigenvalues of the measurement.
             This can only be specified if an observable was not provided.
-        shape (tuple[int] or None): The output shape of the measurement proccess. For
-            some measurement processes a shape is not applicable or the shape
-            may depend on device options, in such cases ``shape=None``. The
-            shape is determined by the return type. For example, assuming a
-            device with ``shots=None``, expectation values and variances define
-            ``shape=(1,)``, whereas probabilities in the qubit define
-            ``shape=(1, 2**num_wires)`` where ``num_wires`` is the number of
-            wires the measurement acts on.
-        numeric_type (type): The expected Python numeric type of the result;
-            either ``int``, ``float`` or ``complex``.
     """
 
     # pylint: disable=too-few-public-methods
@@ -121,8 +111,6 @@ class MeasurementProcess:
         wires=None,
         eigvals=None,
         id=None,
-        shape=None,
-        numeric_type=None,
         log_base=None,
     ):
         self.return_type = return_type
@@ -142,14 +130,6 @@ class MeasurementProcess:
 
             self._eigvals = np.array(eigvals)
 
-        self._shape = shape
-        """tuple[int] or None: The output shape of the measurement proccess. For some a
-        shape is not applicable or the shape may depend on device options, in
-        such cases ``shape=None``."""
-
-        self.numeric_type = numeric_type
-        """tuple[int]: The Python numeric type of the measurement result."""
-
         # TODO: remove the following lines once devices
         # have been refactored to accept and understand recieving
         # measurement processes rather than specific observables.
@@ -166,6 +146,43 @@ class MeasurementProcess:
         # Queue the measurement process
         self.queue()
 
+    @property
+    @functools.lru_cache()
+    def numeric_type(self):
+        """The Python numeric type of the measurement result.
+
+        Returns:
+            type: The output numeric type; ``int``, ``float`` or ``complex``.
+
+        Raises:
+            QuantumFunctionError: the return type of the measurement process is
+                unrecognized and cannot deduce the numeric type
+        """
+        if self.return_type in (Expectation, MutualInfo, Probability, Variance, VnEntropy):
+            return float
+
+        if self.return_type is State:
+            return complex
+
+        if self.return_type is Sample:
+
+            # Note: we only assume an integer numeric type if the observable is a
+            # built-in observable with integer eigenvalues or a tensor product thereof
+            if self.obs is None:
+
+                # Computational basis samples
+                numeric_type = int
+            else:
+                int_eigval_obs = {qml.PauliX, qml.PauliY, qml.PauliZ, qml.Hadamard, qml.Identity}
+                tensor_terms = [self.obs] if not hasattr(self.obs, "obs") else self.obs.obs
+                every_term_standard = all(o.__class__ in int_eigval_obs for o in tensor_terms)
+                numeric_type = int if every_term_standard else float
+            return numeric_type
+
+        raise qml.QuantumFunctionError(
+            f"Cannot deduce the numeric type of the measurement process with unrecognized return_type {self.return_type}."
+        )
+
     @functools.lru_cache()
     def shape(self, device=None):
         """The expected output shape of the MeasurementProcess.
@@ -176,61 +193,86 @@ class MeasurementProcess:
           ``Sample``;
         * The shot vector was defined in the device.
 
+        For example, assuming a device with ``shots=None``, expectation values
+        and variances define ``shape=(1,)``, whereas probabilities in the qubit
+        model define ``shape=(1, 2**num_wires)`` where ``num_wires`` is the
+        number of wires the measurement acts on.
+
+        Note that the shapes for vector-valued return types such as
+        ``Probability`` and ``State`` are adjusted to the output of
+        ``qml.execute`` and may have an extra first element that is squeezed
+        when using QNodes.
+
         Args:
             device (.Device): a PennyLane device to use for determining the
                 shape
 
         Returns:
             tuple: the output shape
+
+        Raises:
+            QuantumFunctionError: the return type of the measurement process is
+                unrecognized and cannot deduce the numeric type
         """
-        if device is None:
-            if self._shape is not None:
-                return self._shape
+        shape = None
 
-            if self.return_type in (Probability, State, Sample):
-                raise MeasurementShapeError(
-                    f"Return type {self.return_type} requires the device argument to be passed to obtain the shape."
-                )
+        # First: prepare the shape for return types that do not require a
+        # device
+        if self.return_type in (Expectation, MutualInfo, Variance, VnEntropy):
+            shape = (1,)
 
-        if device._shot_vector is None:
+        density_matrix_return = self.return_type == State and self.wires
 
-            # Device dependent shapes
-            if self.return_type == Probability:
-                len_wires = len(self.wires)
-                dim = self._get_num_basis_states(len_wires, device)
-                shape = (1, dim)
+        if density_matrix_return:
+            dim = 2 ** len(self.wires)
+            shape = (1, dim, dim)
 
-            elif self.return_type == State and self._shape is None:
+        # Determine shape if device with shot vector
+        if device is not None and device._shot_vector is not None:
+            shape = self._shot_vector_shape(device, main_shape=shape)
 
-                # Note: qml.density_matrix has its shape defined, so we're handling
-                # the qml.state case; acts on all device wires
-                dim = 2 ** len(device.wires)
-                shape = (1, dim)
+        # If we have a shape, return it here
+        if shape is not None:
+            return shape
 
-            elif self.return_type == Sample:
-                len_wires = len(device.wires)
+        # Then: handle return types that require a device; no shot vector
+        if device is None and self.return_type in (Probability, State, Sample):
+            raise MeasurementShapeError(
+                f"The device argument is required to obtain the shape of the measurement process; got return type {self.return_type}."
+            )
 
-                if self.obs is not None:
-                    # qml.sample(some_observable) case
-                    shape = (1, device.shots)
+        if self.return_type == Probability:
+            len_wires = len(self.wires)
+            dim = self._get_num_basis_states(len_wires, device)
+            return (1, dim)
 
-                else:
-                    # qml.sample() case
-                    shape = (1, device.shots, len_wires)
-            else:
-                shape = self._shape
+        if self.return_type == State:
 
-            # The other return types should have their shapes pre-defined
-            return shape
+            # Note: qml.density_matrix has its shape defined, so we're handling
+            # the qml.state case; acts on all device wires
+            dim = 2 ** len(device.wires)
+            return (1, dim)
 
-        return self._shot_vector_shape(device)
+        if self.return_type == Sample:
+            len_wires = len(device.wires)
+
+            if self.obs is not None:
+                # qml.sample(some_observable) case
+                return (1, device.shots)
+
+            # qml.sample() case
+            return (1, device.shots, len_wires)
+
+        raise qml.QuantumFunctionError(
+            f"Cannot deduce the shape of the measurement process with unrecognized return_type {self.return_type}."
+        )
 
     @functools.lru_cache()
-    def _shot_vector_shape(self, device):
+    def _shot_vector_shape(self, device, main_shape=None):
         """Auxiliary function for getting the output shape when the device has
         the shot vector defined.
 
-        The shape is device dependent even if the return type has a shape
+        The shape is device dependent even if the return type has a main shape
         pre-defined (e.g., expectation values, states, etc.).
         """
         shot_vector = device._shot_vector
@@ -238,10 +280,10 @@ class MeasurementProcess:
         num_shot_elements = sum([s.copies for s in shot_vector])
         shape = ()
 
-        if self._shape is not None:
+        if main_shape is not None:
 
             # Expval, var and density_matrix case
-            shape = list(self._shape)
+            shape = list(main_shape)
             shape[0] *= num_shot_elements
             shape = tuple(shape)
 
@@ -500,7 +542,7 @@ def expval(op):
             f"{op.name} is not an observable: cannot be used with expval"
         )
 
-    return MeasurementProcess(Expectation, obs=op, shape=(1,), numeric_type=float)
+    return MeasurementProcess(Expectation, obs=op)
 
 
 def var(op):
@@ -533,7 +575,7 @@ def var(op):
     if not op.is_hermitian:
         raise qml.QuantumFunctionError(f"{op.name} is not an observable: cannot be used with var")
 
-    return MeasurementProcess(Variance, obs=op, shape=(1,), numeric_type=float)
+    return MeasurementProcess(Variance, obs=op)
 
 
 def sample(op=None, wires=None):
@@ -613,20 +655,6 @@ def sample(op=None, wires=None):
             f"{op.name} is not an observable: cannot be used with sample"
         )
 
-    # Note: we only assume an integer numeric type if the observable is a
-    # built-in observable with integer eigenvalues or a tensor product thereof
-    if op is None:
-
-        # Computational basis samples
-        numeric_type = int
-    else:
-        int_eigval_obs = {qml.PauliX, qml.PauliY, qml.PauliZ, qml.Hadamard, qml.Identity}
-        tensor_terms = [op] if not hasattr(op, "obs") else op.obs
-        if not all(o.__class__ in int_eigval_obs for o in tensor_terms):
-            numeric_type = float
-        else:
-            numeric_type = int
-
     if wires is not None:
         if op is not None:
             raise ValueError(
@@ -634,11 +662,9 @@ def sample(op=None, wires=None):
                 "provided. The wires to sample will be determined directly from the observable."
             )
 
-        return MeasurementProcess(
-            Sample, obs=op, wires=qml.wires.Wires(wires), numeric_type=numeric_type
-        )
+        return MeasurementProcess(Sample, obs=op, wires=qml.wires.Wires(wires))
 
-    return MeasurementProcess(Sample, obs=op, numeric_type=numeric_type)
+    return MeasurementProcess(Sample, obs=op)
 
 
 def probs(wires=None, op=None):
@@ -718,18 +744,14 @@ def probs(wires=None, op=None):
             f"{op} does not define diagonalizing gates : cannot be used to rotate the probability"
         )
 
-    numeric_type = float
-
     if wires is not None:
         if op is not None:
             raise qml.QuantumFunctionError(
                 "Cannot specify the wires to probs if an observable is "
                 "provided. The wires for probs will be determined directly from the observable."
             )
-        return MeasurementProcess(
-            Probability, wires=qml.wires.Wires(wires), numeric_type=numeric_type
-        )
-    return MeasurementProcess(Probability, obs=op, numeric_type=numeric_type)
+        return MeasurementProcess(Probability, wires=qml.wires.Wires(wires))
+    return MeasurementProcess(Probability, obs=op)
 
 
 def state():
@@ -786,7 +808,7 @@ def state():
         -0.07471906623679961
     """
     # pylint: disable=protected-access
-    return MeasurementProcess(State, numeric_type=complex)
+    return MeasurementProcess(State)
 
 
 def density_matrix(wires):
@@ -828,9 +850,7 @@ def density_matrix(wires):
     """
     # pylint: disable=protected-access
     wires = qml.wires.Wires(wires)
-    dim = 2 ** len(wires)
-    shape = (1, dim, dim)
-    return MeasurementProcess(State, wires=wires, shape=shape, numeric_type=complex)
+    return MeasurementProcess(State, wires=wires)
 
 
 def vn_entropy(wires, log_base=None):
@@ -866,9 +886,7 @@ def vn_entropy(wires, log_base=None):
         with a compatible device.
     """
     wires = qml.wires.Wires(wires)
-    return MeasurementProcess(
-        VnEntropy, wires=wires, shape=(1,), log_base=log_base, numeric_type=float
-    )
+    return MeasurementProcess(VnEntropy, wires=wires, log_base=log_base)
 
 
 def mutual_info(wires0, wires1, log_base=None):
@@ -923,9 +941,7 @@ def mutual_info(wires0, wires1, log_base=None):
 
     wires0 = qml.wires.Wires(wires0)
     wires1 = qml.wires.Wires(wires1)
-    return MeasurementProcess(
-        MutualInfo, wires=[wires0, wires1], shape=(1,), log_base=log_base, numeric_type=float
-    )
+    return MeasurementProcess(MutualInfo, wires=[wires0, wires1], log_base=log_base)
 
 
 T = TypeVar("T")
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index b71c04979..503241901 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -1159,7 +1159,7 @@ class QuantumTape(AnnotatedQueue):
             for observable in self._measurements:
                 # Note: if one of the sample measurements contains outputs that
                 # are real, then the entire result will be real
-                if observable.numeric_type == float:
+                if observable.numeric_type is float:
                     return observable.numeric_type
 
             return int


Test patch:
diff --git a/pennylane/devices/tests/test_wires.py b/pennylane/devices/tests/test_wires.py
index c85ef2e3d..32a8d343e 100644
--- a/pennylane/devices/tests/test_wires.py
+++ b/pennylane/devices/tests/test_wires.py
@@ -40,6 +40,7 @@ def make_simple_circuit_expval(device, wires):
 # =====
 
 
+# pylint: disable=too-few-public-methods
 class TestWiresIntegration:
     """Test that the device integrates with PennyLane's wire management."""
 
@@ -54,7 +55,9 @@ class TestWiresIntegration:
         ],
     )
     @pytest.mark.parametrize("circuit_factory", [make_simple_circuit_expval])
-    def test_wires_expval(self, device, circuit_factory, wires1, wires2, tol):
+    def test_wires_expval(
+        self, device, circuit_factory, wires1, wires2, tol
+    ):  # pylint: disable=too-many-arguments
         """Test that the expectation of a circuit is independent from the wire labels used."""
         dev1 = device(wires1)
         dev2 = device(wires2)
diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py
index 29e0bdab5..645b01f32 100644
--- a/tests/interfaces/test_jax.py
+++ b/tests/interfaces/test_jax.py
@@ -19,6 +19,11 @@ pytestmark = pytest.mark.jax
 
 jax = pytest.importorskip("jax")
 jnp = pytest.importorskip("jax.numpy")
+
+from jax.config import config
+
+config.update("jax_enable_x64", True)
+
 import numpy as np
 
 import pennylane as qml
@@ -664,43 +669,10 @@ class TestJaxExecuteIntegration:
         res = jax.grad(cost)(params, cache=None)
         assert res.shape == (3,)
 
-    @pytest.mark.parametrize(
-        "ret, mes",
-        [
-            ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], "single return type"),
-            ([qml.state()], "Only Variance and Expectation"),
-        ],
-    )
-    def test_raises_for_jax_jit(self, execute_kwargs, interface, ret, mes):
-        """Tests multiple measurements and unsupported measurements raise an
-        error for the jit JAX interface."""
-        dev = qml.device("default.qubit", wires=2)
-        params = jnp.array([0.1, 0.2, 0.3])
-
-        def cost(a, cache):
-            with qml.tape.QuantumTape() as tape:
-                qml.RY(a[0], wires=0)
-                qml.RX(a[1], wires=0)
-                qml.RY(a[2], wires=0)
-                [qml.apply(r) for r in ret]
-
-            res = qml.interfaces.execute(
-                # Test only applicable for the jax jit interface
-                [tape],
-                dev,
-                cache=cache,
-                interface="jax-jit",
-                **execute_kwargs
-            )
-            return res[0][0]
-
-        with pytest.raises(InterfaceUnsupportedError, match=mes):
-            cost(params, cache=None)
-
 
 @pytest.mark.parametrize("execute_kwargs", execute_kwargs)
 class TestVectorValued:
-    """Test vector-valued returns for the JAX Python interface."""
+    """Test vector-valued jacobian returns for the JAX Python interface."""
 
     def test_multiple_expvals(self, execute_kwargs):
         """Tests computing multiple expectation values in a tape."""
@@ -862,6 +834,231 @@ class TestVectorValued:
             jax.jacobian(cost)(params, cache=None)
 
 
+@pytest.mark.parametrize("execute_kwargs", execute_kwargs)
+class TestVectorValuedJIT:
+    """Test vector-valued returns for the JAX jit Python interface."""
+
+    @pytest.mark.parametrize(
+        "ret_type, shape",
+        [
+            ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (2,)),
+            ([qml.probs(wires=[0, 1])], (1, 4)),
+        ],
+    )
+    def test_shapes(self, execute_kwargs, ret_type, shape):
+        """Test the shape of the result of vector-valued QNodes."""
+        adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian"
+        if adjoint:
+            pytest.skip("The adjoint diff method doesn't support probabilities.")
+
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.1, 0.2, 0.3])
+
+        idx = 0
+
+        def cost(a, cache):
+            with qml.tape.QuantumTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+                for r in ret_type:
+                    qml.apply(r)
+
+            res = qml.interfaces.execute(
+                [tape], dev, cache=cache, interface="jax-jit", **execute_kwargs
+            )
+            return res[0]
+
+        res = cost(params, cache=None)
+        assert res.shape == shape
+
+    def test_independent_expval(self, execute_kwargs):
+        """Tests computing an expectation value that is independent trainable
+        parameters."""
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.1, 0.2, 0.3])
+
+        def cost(a, cache):
+            with qml.tape.QuantumTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+                qml.expval(qml.PauliZ(1))
+
+            res = qml.interfaces.execute(
+                [tape], dev, cache=cache, interface="jax-jit", **execute_kwargs
+            )
+            return res[0][0]
+
+        res = jax.grad(cost)(params, cache=None)
+        assert res.shape == (3,)
+
+    ret_and_output_dim = [
+        ([qml.probs(wires=0)], (1, 2)),
+        ([qml.state()], (1, 4)),
+        ([qml.density_matrix(wires=0)], (1, 2, 2)),
+        # Multi measurements
+        ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (2,)),
+        ([qml.var(qml.PauliZ(0)), qml.var(qml.PauliZ(1))], (2,)),
+        ([qml.probs(wires=0), qml.probs(wires=1)], (2, 2)),
+    ]
+
+    @pytest.mark.parametrize("ret, out_dim", ret_and_output_dim)
+    def test_vector_valued_qnode(self, execute_kwargs, ret, out_dim):
+        """Tests the shape of vector-valued QNode results."""
+
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.1, 0.2, 0.3])
+
+        grad_meth = (
+            execute_kwargs["gradient_kwargs"]["method"]
+            if "gradient_kwargs" in execute_kwargs
+            else ""
+        )
+        if "adjoint" in grad_meth and any(
+            r.return_type
+            in (qml.measurements.Probability, qml.measurements.State, qml.measurements.Variance)
+            for r in ret
+        ):
+            pytest.skip("Adjoint does not support probs")
+
+        def cost(a, cache):
+            with qml.tape.QuantumTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+
+                for r in ret:
+                    qml.apply(r)
+
+            res = qml.interfaces.execute(
+                [tape], dev, cache=cache, interface="jax-jit", **execute_kwargs
+            )[0]
+            return res
+
+        res = cost(params, cache=None)
+        assert res.shape == out_dim
+
+    def test_qnode_sample(self, execute_kwargs):
+        """Tests computing multiple expectation values in a tape."""
+        dev = qml.device("default.qubit", wires=2, shots=10)
+        params = jnp.array([0.1, 0.2, 0.3])
+
+        grad_meth = (
+            execute_kwargs["gradient_kwargs"]["method"]
+            if "gradient_kwargs" in execute_kwargs
+            else ""
+        )
+        if "adjoint" in grad_meth or "backprop" in grad_meth:
+            pytest.skip("Adjoint does not support probs")
+
+        def cost(a, cache):
+            with qml.tape.QuantumTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+                qml.sample(qml.PauliZ(0))
+
+            res = qml.interfaces.execute(
+                [tape], dev, cache=cache, interface="jax-jit", **execute_kwargs
+            )[0]
+            return res
+
+        res = cost(params, cache=None)
+        assert res.shape == (1, dev.shots)
+
+    def test_multiple_expvals_grad(self, execute_kwargs):
+        """Tests computing multiple expectation values in a tape."""
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.1, 0.2, 0.3])
+        fwd_mode = execute_kwargs.get("mode", "not forward") == "forward"
+        if fwd_mode:
+            pytest.skip("The forward mode is tested separately as it should raise an error.")
+
+        def cost(a, cache):
+            with qml.tape.QuantumTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+                qml.expval(qml.PauliZ(0))
+                qml.expval(qml.PauliZ(1))
+
+            res = qml.interfaces.execute(
+                [tape], dev, cache=cache, interface="jax-jit", **execute_kwargs
+            )[0]
+            return res[0] + res[1]
+
+        res = jax.grad(cost)(params, cache=None)
+        assert res.shape == (3,)
+
+    def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs):
+        """Test the jacobian computation with multiple tapes with probability
+        and expectation value computations."""
+        fwd_mode = execute_kwargs.get("mode", "not forward") == "forward"
+        if fwd_mode:
+            pytest.skip("The forward mode is tested separately as it should raise an error.")
+
+        adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian"
+        if adjoint:
+            pytest.skip("The adjoint diff method doesn't support probabilities.")
+
+        def cost(x, y, device, interface, ek):
+            with qml.tape.QuantumTape() as tape1:
+                qml.RX(x, wires=[0])
+                qml.RY(y, wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.expval(qml.PauliZ(0))
+                qml.expval(qml.PauliZ(1))
+
+            with qml.tape.QuantumTape() as tape2:
+                qml.RX(x, wires=[0])
+                qml.RY(y, wires=[1])
+                qml.CNOT(wires=[0, 1])
+                qml.probs(wires=[0])
+                qml.probs(wires=[1])
+
+            return qml.execute([tape1, tape2], device, **ek, interface=interface)[0]
+
+        dev = qml.device("default.qubit", wires=2)
+        x = jnp.array(0.543)
+        y = jnp.array(-0.654)
+
+        x_ = np.array(0.543)
+        y_ = np.array(-0.654)
+
+        res = cost(x, y, dev, interface="jax-jit", ek=execute_kwargs)
+
+        exp = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs)
+
+        for r, e in zip(res, exp):
+            assert jnp.allclose(r, e, atol=1e-7)
+
+    def test_multiple_expvals_raises_fwd_device_grad(self, execute_kwargs):
+        """Tests computing multiple expectation values in a tape."""
+        fwd_mode = execute_kwargs.get("mode", "not forward") == "forward"
+        if not fwd_mode:
+            pytest.skip("Forward mode is not turned on.")
+
+        dev = qml.device("default.qubit", wires=2)
+        params = jnp.array([0.1, 0.2, 0.3])
+
+        def cost(a, cache):
+            with qml.tape.QuantumTape() as tape:
+                qml.RY(a[0], wires=0)
+                qml.RX(a[1], wires=0)
+                qml.RY(a[2], wires=0)
+                qml.expval(qml.PauliZ(0))
+                qml.expval(qml.PauliZ(1))
+
+            res = qml.interfaces.execute(
+                [tape], dev, cache=cache, interface="jax-jit", **execute_kwargs
+            )
+            return res[0]
+
+        with pytest.raises(InterfaceUnsupportedError):
+            jax.jacobian(cost)(params, cache=None)
+
+
 def test_diff_method_None_jit():
     """Test that jitted execution works when `gradient_fn=None`."""
 
diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py
index f46440514..17ae71b2f 100644
--- a/tests/interfaces/test_jax_qnode.py
+++ b/tests/interfaces/test_jax_qnode.py
@@ -221,6 +221,11 @@ vv_qubit_device_and_diff_method = [
     ["default.qubit", "parameter-shift", "backward", "jax-python"],
     ["default.qubit", "adjoint", "forward", "jax-python"],
     ["default.qubit", "adjoint", "backward", "jax-python"],
+    # Jit
+    ["default.qubit", "finite-diff", "backward", "jax-jit"],
+    ["default.qubit", "parameter-shift", "backward", "jax-jit"],
+    ["default.qubit", "adjoint", "forward", "jax-jit"],
+    ["default.qubit", "adjoint", "backward", "jax-jit"],
 ]
 
 
@@ -261,6 +266,11 @@ class TestVectorValuedQNode:
         expected = [np.cos(a), -np.cos(a) * np.sin(b)]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+        if diff_method in {"finite-diff", "parameter-shift"} and interface == "jax-jit":
+
+            # No jax.jacobian support for call
+            pytest.xfail(reason="batching rules are implemented only for id_tap, not for call.")
+
         res = jax.jacobian(circuit, argnums=[0, 1])(a, b)
         expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T
         assert np.allclose(res, expected, atol=tol, rtol=0)
@@ -323,6 +333,10 @@ class TestVectorValuedQNode:
             qml.CNOT(wires=[0, 1])
             return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]
 
+        if diff_method in {"finite-diff", "parameter-shift"} and interface == "jax-jit":
+            # No jax.jacobian support for call
+            pytest.xfail(reason="batching rules are implemented only for id_tap, not for call.")
+
         jac_fn = jax.jacobian(circuit, argnums=[0, 1])
         res = jac_fn(a, b)
         expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T
@@ -356,6 +370,10 @@ class TestVectorValuedQNode:
             qml.RX(a[1], wires=0)
             return qml.expval(qml.PauliZ(0))
 
+        if diff_method in {"finite-diff", "parameter-shift"} and interface == "jax-jit":
+            # No jax.jacobian support for call
+            pytest.xfail(reason="batching rules are implemented only for id_tap, not for call.")
+
         jax.jacobian(circuit)(a)
 
         for args in spy.call_args_list:
@@ -476,11 +494,6 @@ class TestQubitIntegration:
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
 
-        if interface == "jax-jit":
-            pytest.skip(
-                "Only Variance and Expectation returns are supported for the jittable JAX interface."
-            )
-
         dev = qml.device(dev_name, wires=2)
         x = jnp.array(0.543)
         y = jnp.array(-0.654)
@@ -492,6 +505,11 @@ class TestQubitIntegration:
             qml.CNOT(wires=[0, 1])
             return qml.probs(wires=[1])
 
+        if diff_method in {"finite-diff", "parameter-shift"} and interface == "jax-jit":
+
+            # No jax.jacobian support for call
+            pytest.xfail(reason="batching rules are implemented only for id_tap, not for call.")
+
         res = jax.jacobian(circuit, argnums=[0, 1])(x, y)
 
         expected = np.array(
@@ -508,14 +526,16 @@ class TestQubitIntegration:
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support probs")
 
-        if interface == "jax-jit":
-            pytest.skip(
-                "Only Variance and Expectation returns are supported for the jittable JAX interface."
-            )
+        if diff_method in {"finite-diff", "parameter-shift"} and interface == "jax-jit":
+
+            # No jax.jacobian support for call
+            pytest.xfail(reason="batching rules are implemented only for id_tap, not for call.")
 
         dev = qml.device(dev_name, wires=3)
         x = jnp.array(0.543)
         y = jnp.array(-0.654)
+        if diff_method == "adjoint":
+            pytest.skip("Adjoint does not support probs")
 
         @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
         def circuit(x, y):
@@ -686,11 +706,6 @@ class TestQubitIntegration:
         if diff_method == "adjoint":
             pytest.skip("Adjoint warns with finite shots")
 
-        if interface == "jax-jit":
-            pytest.skip(
-                "Only Variance and Expectation returns are supported for the jittable JAX interface."
-            )
-
         dev = qml.device(dev_name, wires=2, shots=10)
 
         @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
@@ -967,11 +982,6 @@ class TestQubitIntegration:
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support states")
 
-        if interface == "jax-jit":
-            pytest.skip(
-                "Only Variance and Expectation returns are supported for the jittable JAX interface."
-            )
-
         dev = qml.device(dev_name, wires=2)
 
         x = jnp.array(0.543)
@@ -1223,11 +1233,6 @@ class TestTapeExpansion:
         if diff_method in ("adjoint", "backprop", "finite-diff"):
             pytest.skip("The adjoint and backprop methods do not yet support sampling")
 
-        if interface == "jax-jit":
-            pytest.skip(
-                "Only Variance and Expectation returns are supported for the jittable JAX interface."
-            )
-
         if max_diff > 1:
             pytest.skip("JAX only supports first derivatives")
 
@@ -1392,3 +1397,43 @@ class TestJIT:
 
         expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)]
         assert np.allclose(g, expected_g, atol=tol, rtol=0)
+
+    def test_gradient_scalar_cost_vector_valued_qnode(self, dev_name, diff_method, mode, tol):
+        """Test derivative calculation of a scalar valued cost function that
+        uses the output of a vector-valued QNode"""
+        dev = qml.device(dev_name, wires=2)
+
+        if diff_method == "adjoint":
+            pytest.xfail(reason="The adjoint method is not using host-callback currently")
+
+        @jax.jit
+        @qnode(dev, diff_method=diff_method, interface="jax", mode=mode)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[1])
+
+        def cost(x, y, idx):
+            res = circuit(x, y)
+            return res[idx]
+
+        x = jnp.array(1.0)
+        y = jnp.array(2.0)
+        expected_g = (
+            np.array([-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2]),
+            np.array([-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]),
+        )
+
+        idx = 0
+        g0 = jax.grad(cost, argnums=0)(x, y, idx)
+        g1 = jax.grad(cost, argnums=1)(x, y, idx)
+        assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0)
+        assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0)
+
+        idx = 1
+        g0 = jax.grad(cost, argnums=0)(x, y, idx)
+        g1 = jax.grad(cost, argnums=1)(x, y, idx)
+
+        assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0)
+        assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0)
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index e1a726ba8..f1372172c 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1299,6 +1299,7 @@ class TestExecution:
         expected = [np.cos(x), np.cos(y) ** 2]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
     def test_prob_expectation_values(self, tol):
         """Tests correct output shape and evaluation for a tape
         with prob and var outputs"""
diff --git a/tests/test_measurements.py b/tests/test_measurements.py
index 3a01ad468..65b02f755 100644
--- a/tests/test_measurements.py
+++ b/tests/test_measurements.py
@@ -75,6 +75,22 @@ def test_no_measure(tol):
         res = circuit(0.65)
 
 
+def test_numeric_type_unrecognized_error(tol):
+    """Test that querying the numeric type of a measurement process with an
+    unrecognized return type raises an error."""
+    mp = MeasurementProcess("NotValidReturnType")
+    with pytest.raises(qml.QuantumFunctionError, match="Cannot deduce the numeric type"):
+        mp.numeric_type()
+
+
+def test_shape_unrecognized_error(tol):
+    """Test that querying the shape of a measurement process with an
+    unrecognized return type raises an error."""
+    mp = MeasurementProcess("NotValidReturnType")
+    with pytest.raises(qml.QuantumFunctionError, match="Cannot deduce the shape"):
+        mp.shape()
+
+
 class TestExpval:
     """Tests for the expval function"""
 
@@ -261,7 +277,7 @@ class TestProbs:
         the shape of certain measurements."""
         with pytest.raises(
             MeasurementShapeError,
-            match="requires the device argument to be passed to obtain the shape",
+            match="The device argument is required to obtain the shape of the measurement process",
         ):
             measurement.shape()
 
@@ -284,6 +300,7 @@ class TestSample:
 
         assert np.array_equal(sample.shape, (2, n_sample))
 
+    @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
     def test_sample_combination(self, tol):
         """Test the output of combining expval, var and sample"""
         n_sample = 10
@@ -338,6 +355,7 @@ class TestSample:
         assert np.array_equal(result.shape, (3, n_sample))
         assert result.dtype == np.dtype("int")
 
+    @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
     def test_sample_output_type_in_combination(self, tol):
         """Test the return type and shape of sampling multiple works
         in combination with expvals and vars"""
@@ -1152,20 +1170,6 @@ class TestState:
         res = qml.state()
         assert res.shape(dev) == (3, 2**3)
 
-    @pytest.mark.parametrize("shots", [None, 1, 10])
-    def test_shape(self, shots):
-        """Test that the shape is correct for qml.density_matrix."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
-        res = qml.density_matrix(wires=[0, 1])
-        assert res.shape(dev) == (1, 2**2, 2**2)
-
-    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
-    def test_shape_shot_vector(self, s_vec):
-        """Test that the shape is correct for qml.density_matrix with the shot vector too."""
-        dev = qml.device("default.qubit", wires=3, shots=s_vec)
-        res = qml.density_matrix(wires=[0, 1])
-        assert res.shape(dev) == (3, 2**2, 2**2)
-
     def test_numeric_type(self):
         """Test that the numeric type of state measurements."""
         assert qml.state().numeric_type == complex
@@ -1470,3 +1474,17 @@ class TestDensityMatrix:
 
         with pytest.raises(qml.QuantumFunctionError, match="custom wire labels"):
             func()
+
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_shape(self, shots):
+        """Test that the shape is correct for qml.density_matrix."""
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+        res = qml.density_matrix(wires=[0, 1])
+        assert res.shape(dev) == (1, 2**2, 2**2)
+
+    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
+    def test_shape_shot_vector(self, s_vec):
+        """Test that the shape is correct for qml.density_matrix with the shot vector too."""
+        dev = qml.device("default.qubit", wires=3, shots=s_vec)
+        res = qml.density_matrix(wires=[0, 1])
+        assert res.shape(dev) == (3, 2**2, 2**2)