def test_map_fields_as_expected(self, fields_to_map):
    """Tests that the fields are mapped are as expected."""
    graph = graphs.GraphsTuple(**self.graph)
    graph = graph.map(lambda v: v + v, fields_to_map)
    for field in graphs.ALL_FIELDS:
      if field in fields_to_map:
        self.assertEqual(field + field, getattr(graph, field))
      else:
        self.assertEqual(field, getattr(graph, field))
Tests that the fields are mapped are as expected.
graph_nets/tests/graphs_test.py