class_docstr: Test that the transform is differentiable
docstr: Tests that the output of the parameter-shift CV transform
  can be differentiated using JAX, yielding second derivatives.
docstr_tok: ["Tests","that","the","output","of","the","parameter-shift","CV","transform","can","be","differentiated","using","JAX",",","yielding","second","derivatives","."]
code: def test_jax(self, tol):
      """Tests that the output of the parameter-shift CV transform
      can be differentiated using JAX, yielding second derivatives."""
      import jax
      from jax import numpy as jnp
  
      dev = qml.device("default.gaussian", wires=1)
      params = jnp.array([0.543, -0.654])
  
      def cost_fn(x):
          with qml.queuing.AnnotatedQueue() as q:
              qml.Squeezing(x[0], 0, wires=0)
              qml.Rotation(x[1], wires=0)
              qml.var(qml.QuadX(wires=[0]))
  
          tape = qml.tape.QuantumScript.from_queue(q)
          tape.trainable_params = {0, 2}
          tapes, fn = qml.gradients.param_shift_cv(tape, dev)
          jac = fn(
              qml.execute(
                  tapes,
                  dev,
                  param_shift_cv,
                  gradient_kwargs={"dev": dev},
                  interface="jax",
              )
          )
          return jac
  
      r, phi = params
      res = cost_fn(params)
      expected = np.array(
          [
              2 * np.exp(2 * r) * np.sin(phi) ** 2
              - 2 * np.exp(-2 * r) * np.cos(phi) ** 2,
              2 * np.sinh(2 * r) * np.sin(2 * phi),
          ]
      )
  
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
      res = jax.jacobian(cost_fn)(params)
  
      expected = np.array(
          [
              4 * np.cosh(2 * r) * np.sin(2 * phi),
              4 * np.cos(2 * phi) * np.sinh(2 * r),
          ]
      )
  
      assert np.allclose(res[1], expected, atol=tol, rtol=0)
  
imports: import unittest.mock
  import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.gradients import param_shift_cv
  from pennylane.gradients.gradient_transform import choose_trainable_params
  from pennylane.gradients.parameter_shift_cv import _grad_method_cv, _find_gradient_methods_cv, _transform_observable
label_desc: Use image node_id 4 for calling the TestParamShiftInterfaces obj's underlying member method code with example usage: obj.test_jax(tol) and returns: jac