def test_dynamic_batch_sizes(self):
    """Checks that all batch sizes are as expected through a GraphNetwork."""
    input_graph = utils_np.data_dicts_to_graphs_tuple(
        [SMALL_GRAPH_1, SMALL_GRAPH_2])
    input_graph = input_graph.map(tf.constant, fields=graphs.ALL_FIELDS)
    model = self._get_model()
    output = model(input_graph)
    for k, v in input_graph._asdict().items():
      self.assertEqual(v.shape[0], getattr(output, k).shape[0])
Checks that all batch sizes are as expected through a GraphNetwork.
graph_nets/tests_tf2/modules_test.py