def test_edge_block_options(self,
                              use_edges,
                              use_receiver_nodes,
                              use_sender_nodes,
                              use_globals):
    """Test for configuring the EdgeBlock options."""
    reducer = tf.math.unsorted_segment_sum
    input_graph = self._get_input_graph()
    edge_model_fn = functools.partial(snt.Linear, output_size=10)
    edge_block_opt = {"use_edges": use_edges,
                      "use_receiver_nodes": use_receiver_nodes,
                      "use_sender_nodes": use_sender_nodes,
                      "use_globals": use_globals}
    # 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}
    # Identity global model
    global_model_fn = lambda: tf.identity
    global_block_opt = {"use_globals": True,
                        "use_nodes": False,
                        "use_edges": False}

    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,
        reducer=reducer)

    output_graph = utils_tf.nest_to_numpy(graph_network(input_graph))

    edge_block = blocks.EdgeBlock(
        edge_model_fn=lambda: graph_network._edge_block._edge_model,
        use_edges=use_edges,
        use_receiver_nodes=use_receiver_nodes,
        use_sender_nodes=use_sender_nodes,
        use_globals=use_globals)

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

    self._assert_all_none_or_all_close(expected_edges,
                                       output_graph.edges)
    self._assert_all_none_or_all_close(expected_nodes,
                                       output_graph.nodes)
    self._assert_all_none_or_all_close(expected_globals,
                                       output_graph.globals)
Test for configuring the EdgeBlock options.
graph_nets/tests_tf2/modules_test.py