def test_dynamic_batch_sizes(self, block_constructor):
    """Checks that all batch sizes are as expected through a GraphNetwork."""
    # Remove all placeholders from here, these are unnecessary in tf2.
    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 = block_constructor(
        functools.partial(snt.nets.MLP, output_sizes=[10]))
    output = model(input_graph)
    actual = utils_tf.nest_to_numpy(output)

    for k, v in input_graph._asdict().items():
      self.assertEqual(v.shape[0], getattr(actual, k).shape[0])
Checks that all batch sizes are as expected through a GraphNetwork.
graph_nets/tests_tf2/blocks_test.py