class_docstr: Tests for the broadcast_expand transform
class_docstr_tok: ["Tests","for","the","broadcast_expand","transform"]
docstr: Test that the expansion works with autograd and is differentiable.
docstr_tok: ["Test","that","the","expansion","works","with","autograd","and","is","differentiable","."]
code: def test_autograd(self, params, size, obs, exp_fn):
      """Test that the expansion works with autograd and is differentiable."""
      params = tuple(pnp.array(p, requires_grad=True) for p in params)
  
      def cost(*params):
          tape = make_tape(*params, obs)
          tapes, fn = qml.transforms.broadcast_expand(tape)
          if len(tape.measurements) > 1:
              return qml.math.stack(
                  fn(qml.execute(tapes, dev, qml.gradients.param_shift))
              )
          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 = qml.jacobian(cost)(*params)
      exp_jac = qml.jacobian(exp_fn)(*params)
  
      if len(obs) > 1 and size == 1:
          new_exp_jac = [qml.math.flatten(el) for el in exp_jac]
          new_jac = [qml.math.flatten(el) for el in jac]
          assert all(
              qml.math.allclose(_jac, e_jac)
              for _jac, e_jac in zip(new_jac, new_exp_jac)
          )
      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 5 for calling the TestBroadcastExpand obj's underlying member method code with example usage: obj.test_autograd(params, size, obs, exp_fn) and returns: fn, qml