class_docstr: Tests for the batch_transform decorator differentiability
class_docstr_tok: ["Tests","for","the","batch_transform","decorator","differentiability"]
docstr: Test that a batch transform is differentiable when using
  jax
docstr_tok: ["Test","that","a","batch","transform","is","differentiable","when","using","jax"]
code: def test_differentiable_jax(self, diff_method):
      """Test that a batch transform is differentiable when using
      jax"""
      import jax
  
      jax.config.update("jax_enable_x64", True)
  
      dev = qml.device("default.qubit", wires=2)
      qnode = qml.QNode(self.circuit, dev, diff_method=diff_method)
  
      def cost(x, weights):
          # pylint: disable=unexpected-keyword-arg
          return self.my_transform(qnode, weights, max_diff=1)(x)
  
      weights_np = np.array([0.1, 0.2], requires_grad=True)
      x_np = np.array(0.543, requires_grad=True)
      weights = jax.numpy.array(weights_np)
      x = jax.numpy.array(x_np)
  
      res = cost(x, weights)
      assert np.allclose(res, self.expval(x, weights))
  
      grad = jax.grad(cost, argnums=[0, 1])(x, weights)
      expected = qml.grad(self.expval)(x_np, weights_np)
      assert len(grad) == len(expected)
      for g, e in zip(grad, expected):
          assert qml.math.allclose(g, e)
  
imports: import functools
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 7 for calling the TestBatchTransformGradients obj's underlying member method code with example usage: obj.test_differentiable_jax(diff_method) and returns: self