def test_unused_field_can_be_none(
      self, use_edges, use_nodes, use_globals, none_field):
    """Checks that computation can handle non-necessary fields left None."""
    input_graph = self._get_input_graph([none_field])
    global_block = blocks.GlobalBlock(
        global_model_fn=self._global_model_fn,
        use_edges=use_edges,
        use_nodes=use_nodes,
        use_globals=use_globals)
    output_graph = global_block(input_graph)

    model_inputs = []
    if use_edges:
      model_inputs.append(
          blocks.EdgesToGlobalsAggregator(
              tf.math.unsorted_segment_sum)(input_graph))
    if use_nodes:
      model_inputs.append(
          blocks.NodesToGlobalsAggregator(
              tf.math.unsorted_segment_sum)(input_graph))
    if use_globals:
      model_inputs.append(input_graph.globals)

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

    actual_globals = output_graph.globals.numpy()
    model_inputs_out = model_inputs

    expected_output_globals = model_inputs_out * self._scale
    self.assertNDArrayNear(expected_output_globals, actual_globals, err=1e-4)
Checks that computation can handle non-necessary fields left None.
graph_nets/tests_tf2/blocks_test.py