class_docstr: Tests for the broadcast_expand transform
class_docstr_tok: ["Tests","for","the","broadcast_expand","transform"]
docstr: Test that the expansion works with jax and is differentiable.
docstr_tok: ["Test","that","the","expansion","works","with","jax","and","is","differentiable","."]
code: def test_jax(self, params, size, obs, exp_fn):
      """Test that the expansion works with jax and is differentiable."""
      import jax
  
      params = tuple(jax.numpy.array(p) for p in params)
  
      def cost(*params):
          tape = make_tape(*params, obs)
          tapes, fn = qml.transforms.broadcast_expand(tape)
          return fn(qml.execute(tapes, dev, qml.gradients.param_shift))
  
      expected = exp_fn(*params)
  
      if len(obs) > 1 and size == 1:
          expected = expected.T
  
      assert qml.math.allclose(cost(*params), expected)
  
      jac = jax.jacobian(cost, argnums=[0, 1, 2])(*params)
  
      exp_jac = jax.jacobian(exp_fn, argnums=[0, 1, 2])(*params)
  
      if len(obs) > 1:
          exp_jac_0 = jax.jacobian(exp_fn_Z0, argnums=[0, 1, 2])(
              *params
          )
          exp_jac_1 = jax.jacobian(exp_fn_Y1, argnums=[0, 1, 2])(
              *params
          )
  
          assert all(
              qml.math.allclose(_jac, e_jac)
              for _jac, e_jac in zip(jac[0], exp_jac_0)
          )
          assert all(
              qml.math.allclose(_jac, e_jac)
              for _jac, e_jac in zip(jac[1], exp_jac_1)
          )
  
      else:
          assert all(
              qml.math.allclose(_jac, e_jac)
              for _jac, e_jac in zip(jac, exp_jac)
          )
  
file_path: pennylane/tests/transforms/test_broadcast_expand.py
imports: import pytest
  import numpy
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 6 for calling the TestBroadcastExpand obj's underlying member method code with example usage: obj.test_jax(params, size, obs, exp_fn) and returns: fn