class_docstr: Tests the application and differentiation of an RX gate in the different interfaces.
class_docstr_tok: ["Tests","the","application","and","differentiation","of","an","RX","gate","in","the","different","interfaces","."]
docstr: Test that the application of an rx gate is differentiable with jax.
docstr_tok: ["Test","that","the","application","of","an","rx","gate","is","differentiable","with","jax","."]
code: def test_rx_grad_jax(self, method, use_jit):
      """Test that the application of an rx gate is differentiable with jax."""
  
      import jax
  
      state = jax.numpy.array(self.state)
  
      def f(phi):
          op = qml.RX(phi, wires=0)
          return method(op, state)
  
      if use_jit:
          f = jax.jit(f)
  
      phi = 0.325
  
      new_state = f(phi)
      g = jax.jacobian(f, holomorphic=True)(phi + 0j)
      self.compare_expected_result(phi, state, new_state, g)
  
imports: from functools import reduce
  import pytest
  import numpy
  from scipy.stats import unitary_group
  import pennylane
  from pennylane.operation import _UNSET_BATCH_SIZE, Operation
  from pennylane.devices.qubit.apply_operation import apply_operation, apply_operation_einsum, apply_operation_tensordot
label_desc: Use image node_id 3 for calling the TestRXCalcGrad obj's underlying member method code with example usage: obj.test_rx_grad_jax(method, use_jit) and returns: method