def test_same_as_subblocks(self, reducer, use_tf_function):
    """Compares the output to explicit subblocks output.

    Args:
      reducer: The reducer used in the `NodeBlock` and `GlobalBlock`.
      use_tf_function: Whether to compile the model with `tf.function`.
    """
    input_graph = self._get_input_graph()

    edge_model_fn = functools.partial(snt.Linear, output_size=5)
    node_model_fn = functools.partial(snt.Linear, output_size=10)
    global_model_fn = functools.partial(snt.Linear, output_size=15)

    graph_network = modules.GraphNetwork(
        edge_model_fn=edge_model_fn,
        node_model_fn=node_model_fn,
        global_model_fn=global_model_fn,
        reducer=reducer)

    if use_tf_function:
      input_signature = [utils_tf.specs_from_graphs_tuple(input_graph)]
      graph_network_fn = tf.function(graph_network, input_signature)
    else:
      graph_network_fn = graph_network

    output_graph = graph_network_fn(input_graph)

    edge_block = blocks.EdgeBlock(
        edge_model_fn=lambda: graph_network._edge_block._edge_model,
        use_sender_nodes=True,
        use_edges=True,
        use_receiver_nodes=True,
        use_globals=True)
    node_block = blocks.NodeBlock(
        node_model_fn=lambda: graph_network._node_block._node_model,
        use_nodes=True,
        use_sent_edges=False,
        use_received_edges=True,
        use_globals=True,
        received_edges_reducer=reducer)
    global_block = blocks.GlobalBlock(
        global_model_fn=lambda: graph_network._global_block._global_model,
        use_nodes=True,
        use_edges=True,
        use_globals=True,
        edges_reducer=reducer,
        nodes_reducer=reducer)

    expected_output_edge_block = edge_block(input_graph)
    expected_output_node_block = node_block(expected_output_edge_block)
    expected_output_global_block = 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.numpy())
    self._assert_all_none_or_all_close(expected_nodes,
                                       output_graph.nodes.numpy())
    self._assert_all_none_or_all_close(expected_globals,
                                       output_graph.globals.numpy())
Compares the output to explicit subblocks output.

    Args:
      reducer: The reducer used in the `NodeBlock` and `GlobalBlock`.
      use_tf_function: Whether to compile the model with `tf.function`.
    
graph_nets/tests_tf2/modules_test.py