def test_map_field_called_only_once(self):
    """Tests that the mapping function is called exactly once per field."""
    graph = graphs.GraphsTuple(**self.graph)
    mapped_fields = []
    def map_fn(v):
      mapped_fields.append(v)
      return v
    graph = graph.map(map_fn, graphs.ALL_FIELDS)
    self.assertCountEqual(mapped_fields, graphs.ALL_FIELDS)
Tests that the mapping function is called exactly once per field.
graph_nets/tests/graphs_test.py