def test_incompatible_higher_rank_inputs_no_raise(self, field_to_reshape):
    """A RelationNetwork does not make assumptions on its inputs shapes."""
    input_graph = self._get_shaped_input_graph()
    edge_model_fn, _, global_model_fn = self._get_shaped_model_fns()
    input_graph = input_graph.map(
        lambda v: tf.transpose(v, [0, 2, 1, 3]), [field_to_reshape])
    network = modules.RelationNetwork(edge_model_fn, global_model_fn)
    self._assert_build_and_run(network, input_graph)
A RelationNetwork does not make assumptions on its inputs shapes.
graph_nets/tests/modules_test.py