def test_created_variables(self, name=None):
    """Verifies variable names and shapes created by a RelationNetwork."""
    name = name if name is not None else "relation_network"
    expected_var_shapes_dict = {
        name + "/edge_block/linear/b:0": [5],
        name + "/edge_block/linear/w:0": [2 + 2, 5],
        name + "/global_block/linear/b:0": [15],
        name + "/global_block/linear/w:0": [5, 15],
    }
    input_graph = utils_tf.data_dicts_to_graphs_tuple([SMALL_GRAPH_1])
    model = self._get_model(name=name)

    model(input_graph)
    variables = model.get_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 RelationNetwork.
graph_nets/tests/modules_test.py