class_docstr: Tests for the jvp function
class_docstr_tok: ["Tests","for","the","jvp","function"]
docstr: Tests correct output shape and evaluation for a tape
  with prob and expval outputs and a single parameter
docstr_tok: ["Tests","correct","output","shape","and","evaluation","for","a","tape","with","prob","and","expval","outputs","and","a","single","parameter"]
code: def test_prob_expval_single_param(self, tol, batch_dim):
      """Tests correct output shape and evaluation for a tape
      with prob and expval outputs and a single parameter"""
      if batch_dim is not None:
          pytest.skip(
              msg="JVP computation of batched tapes is disallowed, see #4462"
          )
      dev = qml.device("default.qubit", wires=2)
      x = (
          0.543
          if batch_dim is None
          else 0.543 * np.arange(1, 1 + batch_dim)
      )
  
      with qml.queuing.AnnotatedQueue() as q:
          qml.RX(x, wires=[0])
          qml.CNOT(wires=[0, 1])
          qml.expval(qml.PauliZ(0))
          qml.probs(wires=0)
  
      tape = qml.tape.QuantumScript.from_queue(q)
      tape.trainable_params = {0}
      tangent = np.array([1.0])
  
      tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
      assert len(tapes) == 2
  
      res = fn(dev.execute(tapes))
      assert isinstance(res, tuple)
      assert len(res) == 2
      assert res[0].shape == () if batch_dim is None else (batch_dim,)
      assert (
          res[1].shape == (2,) if batch_dim is None else (batch_dim, 2)
      )
  
      expected_0 = -np.sin(x)
      assert np.allclose(res[0], expected_0, atol=tol, rtol=0)
  
      # Transpose for batch-dimension to be first if present
      expected_1 = np.array([-np.sin(x) / 2, np.sin(x) / 2]).T
      assert np.allclose(res[1], expected_1, atol=tol, rtol=0)
  
imports: import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.gradients import param_shift
label_desc: Use image node_id 9 for calling the TestJVP obj's underlying member method code with example usage: obj.test_prob_expval_single_param(tol, batch_dim) without return types