def compose_chunk(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    combine decomposed chunk
    """

    def match_pattern(module, node):
        if node.op == "call_function" and node.target in (torch.ops.aten.split.Tensor,):
            div = node.args[1]
            input = node.args[0]
            if isinstance(div, int):
                return (False,)
            if div.target != operator.floordiv:
                return (False,)
            else:
                div_const = div.args[1]
                sub = div.args[0]
                if sub.target != operator.sub:
                    return (False,)
                else:
                    add = sub.args[0]
                    if add.target != operator.add:
                        return (False,)
                    else:
                        add_const = add.args[1]
                        if add_const != div_const:
                            return (False,)
                        symsize = add.args[0]
                        if symsize.target != torch.ops.aten.sym_size:
                            return (False,)
                        else:
                            symsize_input = symsize.args[0]
                            dim = symsize.args[1]
                            if symsize_input != input:
                                return (False,)

            return (True, div_const, dim)
        else:
            return (False,)

    modified = False
    for node in module.graph.nodes:
        res = match_pattern(module, node)
        if res[0]:
            modified = True
            with module.graph.inserting_after(node):
                new_args = (node.args[0], res[1], res[2])
                new_node = module.graph.create_node(
                    "call_function",
                    aten_compose_chunk,
                    args=new_args,
                    kwargs=None,
                )
            node.replace_all_uses_with(new_node)

    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)
combine decomposed chunk
fx2ait/fx2ait/passes/lower_basic_pass_aten.py