def test_data_dicts_to_graphs_tuple_cast_types(self):
    """Index and number fields should be cast to numpy arrays."""
    for graph_dict in self.graphs_dicts_in:
      graph_dict["n_node"] = np.array(
          graph_dict["nodes"].shape[0], dtype=np.int64)
      graph_dict["receivers"] = graph_dict["receivers"].astype(np.int16)
      graph_dict["senders"] = graph_dict["senders"].astype(np.float64)
      graph_dict["nodes"] = graph_dict["nodes"].astype(np.float64)
      graph_dict["edges"] = graph_dict["edges"].astype(np.float64)
    out = utils_np.data_dicts_to_graphs_tuple(self.graphs_dicts_in)
    for key in ["n_node", "n_edge", "receivers", "senders"]:
      self.assertEqual(np.int32, getattr(out, key).dtype)
    for key in ["nodes", "edges"]:
      self.assertEqual(np.float64, getattr(out, key).dtype)
Index and number fields should be cast to numpy arrays.
graph_nets/tests/utils_np_test.py