def test_optional_arguments(self, scale, offset):
    """Assesses the correctness of the EdgeBlock using arguments."""
    input_graph = self._get_input_graph()
    edge_block = blocks.EdgeBlock(edge_model_fn=self._edge_model_args_fn)
    output_graph_out = edge_block(
        input_graph, edge_model_kwargs=dict(scale=scale, offset=offset))

    fixed_scale = scale
    fixed_offset = offset
    model_fn = lambda: lambda features: features * fixed_scale + fixed_offset
    hardcoded_edge_block = blocks.EdgeBlock(edge_model_fn=model_fn)
    expected_graph_out = hardcoded_edge_block(input_graph)

    self.assertIs(expected_graph_out.nodes, output_graph_out.nodes)
    self.assertIs(expected_graph_out.globals, output_graph_out.globals)
    self.assertNDArrayNear(
        expected_graph_out.edges.numpy(),
        output_graph_out.edges.numpy(),
        err=1e-4)
Assesses the correctness of the EdgeBlock using arguments.
graph_nets/tests_tf2/blocks_test.py