class_docstr: Test that the jax execute function can be differentiated
class_docstr_tok: ["Test","that","the","jax","execute","function","can","be","differentiated"]
docstr: Tests that the output of the parameter-shift transform
  can be differentiated using jax, yielding second derivatives.
docstr_tok: ["Tests","that","the","output","of","the","parameter-shift","transform","can","be","differentiated","using","jax",",","yielding","second","derivatives","."]
code: def test_parameter_shift_hessian(self, params, tol):
      """Tests that the output of the parameter-shift transform
      can be differentiated using jax, yielding second derivatives."""
      dev = DefaultQubit()
  
      def cost_fn(x):
          ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
          tape1 = qml.tape.QuantumScript(
              ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]
          )
  
          ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))]
          tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)])
          result = execute(
              [tape1, tape2], dev, gradient_fn=param_shift, max_diff=2
          )
          return result[0] + result[1][0]
  
      res = cost_fn(params)
      x, y = params
      expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y))
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
      res = jax.grad(cost_fn)(params)
      expected = jnp.array(
          [
              -jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x),
              -jnp.cos(x) ** 2 * jnp.sin(2 * y),
          ]
      )
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
      res = jax.jacobian(jax.grad(cost_fn))(params)
      expected = jnp.array(
          [
              [
                  -jnp.cos(2 * x) * jnp.cos(2 * y),
                  jnp.sin(2 * x) * jnp.sin(2 * y),
              ],
              [
                  jnp.sin(2 * x) * jnp.sin(2 * y),
                  -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y),
              ],
          ]
      )
      assert np.allclose(res, expected, atol=tol, rtol=0)
  
imports: import pytest
  import numpy
  from param_shift_dev import ParamShiftDerivativesDevice
  import pennylane
  from pennylane.devices import DefaultQubit
  from pennylane.gradients import param_shift
  from pennylane import execute
  from pennylane.measurements import Shots
label_desc: Use image node_id 1 for calling the TestHigherOrderDerivatives obj's underlying member method code with example usage: obj.test_parameter_shift_hessian(params, tol) and returns: unknown