def test_output_values(
      self, use_received_edges, use_sent_edges, use_nodes,
      use_globals, received_edges_reducer, sent_edges_reducer):
    """Compares the output of a NodeBlock to an explicit computation."""
    input_graph = self._get_input_graph()
    node_block = blocks.NodeBlock(
        node_model_fn=self._node_model_fn,
        use_received_edges=use_received_edges,
        use_sent_edges=use_sent_edges,
        use_nodes=use_nodes,
        use_globals=use_globals,
        received_edges_reducer=received_edges_reducer,
        sent_edges_reducer=sent_edges_reducer)
    output_graph = node_block(input_graph)

    model_inputs = []
    if use_received_edges:
      model_inputs.append(
          blocks.ReceivedEdgesToNodesAggregator(
              received_edges_reducer)(input_graph))
    if use_sent_edges:
      model_inputs.append(
          blocks.SentEdgesToNodesAggregator(sent_edges_reducer)(input_graph))
    if use_nodes:
      model_inputs.append(input_graph.nodes)
    if use_globals:
      model_inputs.append(blocks.broadcast_globals_to_nodes(input_graph))

    model_inputs = tf.concat(model_inputs, axis=-1)
    self.assertIs(input_graph.edges, output_graph.edges)
    self.assertIs(input_graph.globals, output_graph.globals)

    output_graph_out = utils_tf.nest_to_numpy(output_graph)
    model_inputs_out = model_inputs

    expected_output_nodes = model_inputs_out * self._scale
    self.assertNDArrayNear(
        expected_output_nodes, output_graph_out.nodes, err=1e-4)
Compares the output of a NodeBlock to an explicit computation.
graph_nets/tests_tf2/blocks_test.py