def _assert_graph_equals_np(self, graph0, graph, force_edges_ordering=False):
    """Asserts that all the graph fields of graph0 and graph match."""

    def silent_convert_to_numpy(tensor):
      if isinstance(tensor, tf.Tensor):
        return tensor.numpy()
      else:
        return tensor

    graph = graph.map(silent_convert_to_numpy, graphs.ALL_FIELDS)
    graph0 = graph0.map(silent_convert_to_numpy, graphs.ALL_FIELDS)

    if graph0.nodes is None:
      self.assertEqual(None, graph.nodes)
    else:
      self.assertAllClose(graph0.nodes, graph.nodes)
    if graph0.globals is None:
      self.assertEqual(None, graph.globals)
    else:
      self.assertAllClose(graph0.globals, graph.globals)
    self.assertAllClose(graph0.n_node, graph.n_node.tolist())
    if graph0.receivers is None:
      self.assertEqual(None, graph.receivers)
      self.assertEqual(None, graph.senders)
      self.assertEqual(None, graph.edges)
      self.assertAllEqual(graph0.n_edge, graph.n_edge)
      return
    self.assertAllClose(graph0.n_edge, graph.n_edge.tolist())

    if not force_edges_ordering:
      self.assertAllClose(graph0.receivers, graph.receivers)
      self.assertAllClose(graph0.senders, graph.senders)
      if graph0.edges is not None:
        self.assertAllClose(graph0.edges, graph.edges)
      else:
        self.assertEqual(None, graph.edges)
      return
    # To compare edges content, we need to make sure they appear in the same
    # order
    if graph0.edges is not None:
      sorted_receivers0, sorted_senders0, sorted_content0 = zip(
          *sorted(zip(graph0.receivers, graph0.senders, graph0.edges.tolist())))
      sorted_receivers, sorted_senders, sorted_content = zip(
          *sorted(zip(graph.receivers, graph.senders, graph.edges.tolist())))
      self.assertAllClose(sorted_content0, sorted_content)
    elif graph.receivers is not None:
      sorted_receivers0, sorted_senders0 = zip(
          *sorted(zip(graph0.receivers, graph0.senders)))
      sorted_receivers, sorted_senders = zip(
          *sorted(zip(graph.receivers, graph.senders)))
    else:
      return
    self.assertAllClose(sorted_receivers0, sorted_receivers)
    self.assertAllClose(sorted_senders0, sorted_senders)
Asserts that all the graph fields of graph0 and graph match.
graph_nets/tests_tf2/test_utils.py