def test_kwargs(self,
                  edge_model_kwargs,
                  node_model_kwargs,
                  global_model_kwargs):
    """Compares the output to expected output graph using kwargs."""
    input_graph = self._get_input_graph()

    edge_model_fn = functools.partial(
        snt.LayerNorm, axis=1, create_scale=False, create_offset=False)
    edge_model_fn_with_params = functools.partial(
        edge_model_fn(),
        scale=edge_model_kwargs["scale"],
        offset=edge_model_kwargs["offset"])

    node_model_fn = functools.partial(
        snt.LayerNorm, axis=1, create_scale=False, create_offset=False)
    node_model_fn_with_params = functools.partial(
        node_model_fn(),
        scale=node_model_kwargs["scale"],
        offset=node_model_kwargs["offset"])

    global_model_fn = functools.partial(
        snt.LayerNorm, axis=1, create_scale=False, create_offset=False)
    global_model_fn_with_params = functools.partial(
        global_model_fn(),
        scale=global_model_kwargs["scale"],
        offset=global_model_kwargs["offset"])

    model = modules.GraphIndependent(
        edge_model_fn, node_model_fn, global_model_fn)
    model_with_params = modules.GraphIndependent(
        lambda: edge_model_fn_with_params,
        lambda: node_model_fn_with_params,
        lambda: global_model_fn_with_params)

    output_model = model(
        input_graph, edge_model_kwargs, node_model_kwargs, global_model_kwargs)
    output_graph = utils_tf.nest_to_numpy(output_model)
    expected_graph = utils_tf.nest_to_numpy(model_with_params(input_graph))

    self.assertAllEqual(expected_graph.receivers, output_graph.receivers,)
    self.assertAllEqual(expected_graph.senders, output_graph.senders)

    self._assert_all_none_or_all_close(expected_graph.edges, output_graph.edges)
    self._assert_all_none_or_all_close(expected_graph.nodes, output_graph.nodes)
    self._assert_all_none_or_all_close(
        expected_graph.globals, output_graph.globals)
Compares the output to expected output graph using kwargs.
graph_nets/tests_tf2/modules_test.py