def test_global_block_options(self,
                                use_edges,
                                use_nodes,
                                use_globals,
                                edges_reducer,
                                nodes_reducer):
    """Test for configuring the NodeBlock options."""
    input_graph = self._get_input_graph()

    if use_edges:
      edges_reducer = edges_reducer or tf.unsorted_segment_sum
    if use_nodes:
      nodes_reducer = nodes_reducer or tf.unsorted_segment_sum

    # Identity edge model.
    edge_model_fn = lambda: tf.identity
    edge_block_opt = {"use_edges": True,
                      "use_receiver_nodes": False,
                      "use_sender_nodes": False,
                      "use_globals": False}
    # Identity node model
    node_model_fn = lambda: tf.identity
    node_block_opt = {"use_received_edges": False,
                      "use_sent_edges": False,
                      "use_nodes": True,
                      "use_globals": False}
    global_model_fn = functools.partial(snt.Linear, output_size=10)
    global_block_opt = {"use_globals": use_globals,
                        "use_nodes": use_nodes,
                        "use_edges": use_edges,
                        "edges_reducer": edges_reducer,
                        "nodes_reducer": nodes_reducer}

    graph_network = modules.GraphNetwork(
        edge_model_fn=edge_model_fn,
        edge_block_opt=edge_block_opt,
        node_model_fn=node_model_fn,
        node_block_opt=node_block_opt,
        global_model_fn=global_model_fn,
        global_block_opt=global_block_opt)

    output_graph = graph_network(input_graph)

    global_block = blocks.GlobalBlock(
        global_model_fn=lambda: graph_network._global_block._global_model,
        use_edges=use_edges,
        use_nodes=use_nodes,
        use_globals=use_globals,
        edges_reducer=edges_reducer,
        nodes_reducer=nodes_reducer)

    expected_output_edge_block = input_graph
    expected_output_node_block = expected_output_edge_block
    expected_output_global_block = global_block(expected_output_node_block)
    expected_edges = expected_output_edge_block.edges
    expected_nodes = expected_output_node_block.nodes
    expected_globals = expected_output_global_block.globals

    with tf.Session() as sess:
      sess.run(tf.global_variables_initializer())
      (output_graph_out,
       expected_edges_out, expected_nodes_out, expected_globals_out) = sess.run(
           (output_graph, expected_edges, expected_nodes, expected_globals))

    self._assert_all_none_or_all_close(expected_edges_out,
                                       output_graph_out.edges)
    self._assert_all_none_or_all_close(expected_nodes_out,
                                       output_graph_out.nodes)
    self._assert_all_none_or_all_close(expected_globals_out,
                                       output_graph_out.globals)
Test for configuring the NodeBlock options.
graph_nets/tests/modules_test.py