def test_incompatible_higher_rank_inputs_no_raise(self):
    """A DeepSets does not make assumptions on the shape if its input edges."""
    input_graph = self._get_shaped_input_graph()
    _, node_model_fn, global_model_fn = self._get_shaped_model_fns()
    input_graph = input_graph.replace(
        edges=tf.transpose(input_graph.edges, [0, 2, 1, 3]))
    network = modules.DeepSets(node_model_fn, global_model_fn)
    self._assert_build_and_run(network, input_graph)
A DeepSets does not make assumptions on the shape if its input edges.
graph_nets/tests_tf2/modules_test.py