def test_incompatible_higher_rank_partial_outputs_no_raise(self):
    """There is no constraint on the size of the partial outputs."""
    input_graph = self._get_shaped_input_graph()
    node_model_fn = functools.partial(
        snt.Conv2D, output_channels=10, kernel_shape=[3, 3], stride=[1, 2])
    global_model_fn = functools.partial(
        snt.Conv2D, output_channels=10, kernel_shape=[3, 3])
    network = modules.DeepSets(node_model_fn, global_model_fn)
    self._assert_build_and_run(network, input_graph)
There is no constraint on the size of the partial outputs.
graph_nets/tests_tf2/modules_test.py