def test_output(self, none_fields):
    """Tests that this function produces the identity."""
    graph_id = self._graph.map(lambda _: None, none_fields)
    graph = graph_id.map(tf.no_op, none_fields)
    expected_out = graph
    actual_out = graph_id
    for field in [
        "nodes", "edges", "globals", "receivers", "senders", "n_node", "n_edge"
    ]:
      if field in none_fields:
        self.assertIsNone(getattr(actual_out, field))
      else:
        self.assertNDArrayNear(
            getattr(expected_out, field), getattr(actual_out, field), err=1e-4)
Tests that this function produces the identity.
graph_nets/tests_tf2/utils_tf_test.py