def test_created_variables(self, name=None):
    """Verifies variable names and shapes created by a GraphNetwork."""
    name = name if name is not None else "graph_network"
    expected_var_shapes_dict = {
        name + "/edge_block/mlp/linear_0/b:0": [5],
        name + "/edge_block/mlp/linear_0/w:0": [4 + 4 + 3, 5],
        name + "/node_block/mlp/linear_0/b:0": [10],
        name + "/node_block/mlp/linear_0/w:0": [5 + 2 + 3, 10],
        name + "/global_block/mlp/linear_0/b:0": [15],
        name + "/global_block/mlp/linear_0/w:0": [10 + 5 + 3, 15],
    }
    input_graph = self._get_input_graph()
    extra_kwargs = {"name": name} if name else {}
    model = modules.GraphNetwork(
        edge_model_fn=functools.partial(snt.nets.MLP, output_sizes=[5]),
        node_model_fn=functools.partial(snt.nets.MLP, output_sizes=[10]),
        global_model_fn=functools.partial(snt.nets.MLP, output_sizes=[15]),
        **extra_kwargs)

    model(input_graph)
    variables = model.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 variable names and shapes created by a GraphNetwork.
graph_nets/tests_tf2/modules_test.py