def test_created_variables(self, use_edges, use_nodes,
                             use_globals, expected_first_dim_w):
    """Verifies the variable names and shapes created by a GlobalBlock."""
    output_size = 10
    expected_var_shapes_dict = {
        "global_block/mlp/linear_0/b:0": [output_size],
        "global_block/mlp/linear_0/w:0": [expected_first_dim_w, output_size]}

    input_graph = self._get_input_graph()

    global_block = blocks.GlobalBlock(
        global_model_fn=functools.partial(snt.nets.MLP,
                                          output_sizes=[output_size]),
        use_edges=use_edges,
        use_nodes=use_nodes,
        use_globals=use_globals)

    global_block(input_graph)

    variables = global_block.variables
    var_shapes_dict = {var.name: var.get_shape().as_list() for var in variables}
    self.assertDictEqual(expected_var_shapes_dict, var_shapes_dict)
Verifies the variable names and shapes created by a GlobalBlock.
graph_nets/tests_tf2/blocks_test.py