def test_higher_rank_outputs(self):
    """Tests that a graph net can be build with higher rank inputs/outputs."""
    input_graph = self._get_shaped_input_graph()
    network = modules.GraphNetwork(*self._get_shaped_model_fns())
    self._assert_build_and_run(network, input_graph)
Tests that a graph net can be build with higher rank inputs/outputs.
graph_nets/tests/modules_test.py