class_docstr: Test vector-valued returns for the JAX-JIT interface.
class_docstr_tok: ["Test","vector-valued","returns","for","the","JAX-JIT","interface","."]
docstr: Test the shape of the result of vector-valued QNodes.
docstr_tok: ["Test","the","shape","of","the","result","of","vector-valued","QNodes","."]
code: def test_shapes(self, execute_kwargs, ret_type, shape, expected_type):
      """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.legacy", wires=2)
      params = jax.numpy.array([0.1, 0.2, 0.3])
  
      def cost(a, cache):
          with qml.queuing.AnnotatedQueue() as q:
              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)
  
          tape = qml.tape.QuantumScript.from_queue(q)
  
          res = qml.execute([tape], dev, cache=cache, **execute_kwargs)
          return res[0]
  
      res = jax.jit(cost)(params, cache=None)
      assert isinstance(res, expected_type)
  
      if expected_type is tuple:
          for r in res:
              assert r.shape == shape
      else:
          assert res.shape == shape
  
imports: import pytest
  import numpy
  import pennylane
  from pennylane.gradients import param_shift
  from pennylane import execute
label_desc: Use image node_id 1 for calling the TestVectorValuedJIT obj's underlying member method code with example usage: obj.test_shapes(execute_kwargs, ret_type, shape, expected_type) and returns: res