def test_gradient_flow(self, use_autograd):
    """Verifies that gradient flow is as expected."""
    input_graph = self._get_input_graph()

    for input_field in ["nodes", "edges", "globals"]:
      input_tensor = getattr(input_graph, input_field)
      for output_field in ["nodes", "edges", "globals"]:
        model = self._get_model()
        if use_autograd:
          model.__call__ = tf.function(model.__call__)
        with tf.GradientTape() as tape:
          tape.watch(input_tensor)
          output_graph = model(input_graph)
        output_tensor = getattr(output_graph, output_field)
        gradient = tape.gradient(output_tensor, input_tensor)
        if input_field == output_field:
          self.assertNotEqual(
              None, gradient,
              msg="gradient should flow from {} to {}".format(
                  output_field, input_field))
        else:
          self.assertIsNone(
              gradient, msg="gradient should not flow from {} to {}".format(
                  output_field, input_field))
Verifies that gradient flow is as expected.
graph_nets/tests_tf2/modules_test.py