code: def test_partial_evaluation_tf(diff_method):
      """Test gradient of partial evaluation matches gradients of
      individual full evaluations using TF"""
      import tensorflow as tf
  
      dev = qml.device("default.qubit", wires=2)
  
      @qml.qnode(dev, diff_method=diff_method)
      def circuit(x, y):
          qml.RX(x, wires=0)
          qml.RY(y[..., 0], wires=0)
          qml.RY(y[..., 1], wires=1)
          return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1))
  
      batch_size = 4
  
      # the partial argument to construct a new circuit with
      y = tf.Variable(np.random.uniform(size=2), trainable=True)
  
      # the batched argument to the new partial circuit
      x = tf.Variable(
          np.random.uniform(size=batch_size), trainable=True
      )
  
      batched_partial_circuit = qml.batch_partial(circuit, y=y)
  
      with tf.GradientTape() as tape:
          out = batched_partial_circuit(x)
  
      # we could also sum over the batch dimension and use the regular
      # gradient instead of the jacobian, but either works
      grad = tf.linalg.tensor_diag_part(tape.jacobian(out, x))
  
      # check the results against individually executed circuits
      indiv_grad = []
      for x_indiv in tf.unstack(x):
          # create a new variable since tensors created by
          # indexing a trainable variable aren't themselves trainable
          x_indiv = tf.Variable(x_indiv, trainable=True)
  
          with tf.GradientTape() as tape:
              out_indiv = circuit(x_indiv, y)
          indiv_grad.append(tape.gradient(out_indiv, x_indiv))
  
      assert np.allclose(grad, indiv_grad)
  
imports: import re
  import pytest
  import pennylane
  from pennylane import numpy
label_desc: Use image node_id 9 for calling a global function with example usage: test_partial_evaluation_tf(diff_method) and returns: qml