def test_incompatible_higher_rank_partial_outputs_raises(self):
    """A error should be raised if partial outputs have incompatible shapes."""
    input_graph = self._get_shaped_input_graph()
    edge_model_fn, node_model_fn, global_model_fn = self._get_shaped_model_fns()
    edge_model_fn_2 = functools.partial(
        snt.Conv2D, output_channels=10, kernel_shape=[3, 3], stride=[1, 2])
    graph_network = modules.GraphNetwork(
        edge_model_fn_2, node_model_fn, global_model_fn)
    with self.assertRaisesRegex(
        tf.errors.InvalidArgumentError, "Dimensions of inputs should match"):
      graph_network(input_graph)
    node_model_fn_2 = functools.partial(
        snt.Conv2D, output_channels=10, kernel_shape=[3, 3], stride=[1, 2])
    graph_network = modules.GraphNetwork(
        edge_model_fn, node_model_fn_2, global_model_fn)
    with self.assertRaisesRegex(
        tf.errors.InvalidArgumentError, "Dimensions of inputs should match"):
      graph_network(input_graph)
A error should be raised if partial outputs have incompatible shapes.
graph_nets/tests_tf2/modules_test.py