def test_incompatible_higher_rank_inputs_raises(self, field_to_reshape):
    """Am exception should be raised if the inputs have incompatible shapes."""
    input_graph = self._get_shaped_input_graph()
    edge_model_fn, node_model_fn, _ = self._get_shaped_model_fns()
    input_graph = input_graph.map(
        lambda v: tf.transpose(v, [0, 2, 1, 3]), [field_to_reshape])
    graph_network = modules.InteractionNetwork(edge_model_fn, node_model_fn)
    with self.assertRaisesRegex(
        tf.errors.InvalidArgumentError, "Dimensions of inputs should match"):
      graph_network(input_graph)
Am exception should be raised if the inputs have incompatible shapes.
graph_nets/tests_tf2/modules_test.py