def test_dtypes(self, data_dtype, indices_dtype, block_constructor):
    """Checks that all the output types are as expected for blocks."""
    input_graph = self._get_input_graph()
    input_graph = input_graph.map(lambda v: tf.cast(v, data_dtype),
                                  ["nodes", "edges", "globals"])
    input_graph = input_graph.map(lambda v: tf.cast(v, indices_dtype),
                                  ["receivers", "senders"])
    model = block_constructor(
        functools.partial(snt.nets.MLP, output_sizes=[10]))
    output = model(input_graph)
    for field in ["nodes", "globals", "edges"]:
      self.assertEqual(data_dtype, getattr(output, field).dtype)
    for field in ["receivers", "senders"]:
      self.assertEqual(indices_dtype, getattr(output, field).dtype)
Checks that all the output types are as expected for blocks.
graph_nets/tests_tf2/blocks_test.py