class_docstr: Tests for the broadcast_expand transform
class_docstr_tok: ["Tests","for","the","broadcast_expand","transform"]
docstr: Test that the expansion works with TensorFlow and is differentiable.
docstr_tok: ["Test","that","the","expansion","works","with","TensorFlow","and","is","differentiable","."]
code: def test_tf(self, params, size, obs, exp_fn):
      """Test that the expansion works with TensorFlow and is differentiable."""
      import tensorflow as tf
  
      params = tuple(tf.Variable(p, dtype=tf.float64) 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))
  
      with tf.GradientTape(persistent=True) as t:
          out = tf.stack(cost(*params))
          exp = exp_fn(*params)
  
      jac = t.jacobian(out, params)
      exp_jac = t.jacobian(exp, 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:
          for _jac, e_jac in zip(jac, exp_jac):
              if e_jac is None:
                  assert qml.math.allclose(_jac, 0.0)
              else:
                  assert qml.math.allclose(_jac, e_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 7 for calling the TestBroadcastExpand obj's underlying member method code with example usage: obj.test_tf(params, size, obs, exp_fn) and returns: fn