def test_data_dicts_to_graphs_tuple_cast_types(self):
    """Index and number fields should be cast to tensors of the right type."""
    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"] = tf.constant(graph_dict["edges"], dtype=tf.float64)
    out = utils_tf.data_dicts_to_graphs_tuple(self.graphs_dicts_in)
    for key in ["n_node", "n_edge", "receivers", "senders"]:
      self.assertEqual(tf.int32, getattr(out, key).dtype)
      self.assertEqual(type(tf.int32), type(getattr(out, key).dtype))
    for key in ["nodes", "edges"]:
      self.assertEqual(type(tf.float64), type(getattr(out, key).dtype))
      self.assertEqual(tf.float64, getattr(out, key).dtype)
Index and number fields should be cast to tensors of the right type.
graph_nets/tests/utils_tf_test.py