class_docstr: Tests for the sum_expand transform
docstr: Tests that the sum_expand tape transform is differentiable with the Jax interface
docstr_tok: ["Tests","that","the","sum_expand","tape","transform","is","differentiable","with","the","Jax","interface"]
code: def test_sum_dif_jax(self, tol):
      """Tests that the sum_expand tape transform is differentiable with the Jax interface"""
      import jax
      from jax import numpy as jnp
  
      S = qml.sum(
          qml.s_prod(-0.2, qml.PauliX(1)),
          qml.s_prod(0.5, qml.prod(qml.PauliZ(1), qml.PauliY(2))),
          qml.s_prod(1, qml.PauliZ(0)),
      )
  
      var = jnp.array([0.1, 0.67, 0.3, 0.4, -0.5, 0.7, -0.2, 0.5, 1])
      output = 0.42294409781940356
      output2 = [
          9.68883500e-02,
          -2.90832724e-01,
          -1.04448033e-01,
          -1.94289029e-09,
          3.50307411e-01,
          -3.41123470e-01,
          0.0,
          -4.36578753e-01,
          6.41233474e-01,
      ]
  
      with AnnotatedQueue() as q:
          for _ in range(2):
              qml.RX(np.array(0), wires=0)
              qml.RX(np.array(0), wires=1)
              qml.RX(np.array(0), wires=2)
              qml.CNOT(wires=[0, 1])
              qml.CNOT(wires=[1, 2])
              qml.CNOT(wires=[2, 0])
  
          qml.expval(S)
  
      qscript = QuantumScript.from_queue(q)
  
      def cost(x):
          new_qscript = qscript.bind_new_parameters(x, list(range(9)))
          tapes, fn = sum_expand(new_qscript)
          res = qml.execute(tapes, dev, qml.gradients.param_shift)
          return fn(res)
  
      assert np.isclose(cost(var), output)
  
      grad = jax.grad(cost)(var)
      assert len(grad) == len(output2)
      for g, o in zip(grad, output2):
          assert np.allclose(g, o, atol=tol)
  
imports: import numpy
  import pytest
  import pennylane
  from pennylane import numpy
  from pennylane.queuing import AnnotatedQueue
  from pennylane.tape import QuantumScript
  from pennylane.transforms import hamiltonian_expand, sum_expand
label_desc: Use image node_id 12 for calling the TestSumExpand obj's underlying member method code with example usage: obj.test_sum_dif_jax(tol) and returns: fn