class_docstr: Test gradients for multiple measurements
class_docstr_tok: ["Test","gradients","for","multiple","measurements"]
docstr: Test derivatives when using JAX
docstr_tok: ["Test","derivatives","when","using","JAX"]
code: def test_jax(self, diff_method, tol, interface):
      """Test derivatives when using JAX"""
      import jax
      import jax.numpy as jnp
  
      jax.config.update("jax_enable_x64", True)
      dev = qml.device("default.qubit", wires=2)
  
      @functools.partial(qml.batch_params, all_operations=True)
      @qml.qnode(dev, diff_method=diff_method, interface=interface)
      def circuit(x):
          qml.RY(x, wires=0)
          qml.CNOT(wires=[0, 1])
          return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
  
      batch_size = 3
      x = jnp.linspace(0.1, 0.5, batch_size)
  
      res = circuit(x)
      expected = (
          jnp.cos(x),
          qml.math.transpose(
              qml.math.stack(
                  [
                      jnp.cos(x / 2) ** 2,
                      jnp.zeros_like(x),
                      jnp.zeros_like(x),
                      jnp.sin(x / 2) ** 2,
                  ]
              )
          ),
      )
  
      assert isinstance(res, tuple)
      assert len(res) == 2
      for r, exp in zip(res, expected):
          assert qml.math.allclose(r, exp, atol=tol)
  
      grad = jax.jacobian(circuit)(x)
      expected = (
          -np.sin(x) * np.eye(batch_size),
          qml.math.stack(
              [
                  -np.sin(x) / 2,
                  np.zeros_like(x),
                  np.zeros_like(x),
                  np.sin(x) / 2,
              ]
          )
          * qml.math.expand_dims(np.eye(batch_size), 1),
      )
  
      assert isinstance(grad, tuple)
      assert len(grad) == 2
      for g, exp in zip(grad, expected):
          assert qml.math.allclose(g, exp, atol=tol, rtol=0)
  
imports: import functools
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 2 for calling the TestDiffMulti obj's underlying member method code with example usage: obj.test_jax(diff_method, tol, interface) and returns: qml,  qml