def _fuse_strided_op_and_cat(sorted_graph: List[Tensor]) -> List[Tensor]:  # noqa: C901
    """
    Fuse strided_ops and cat op. One special case is group_gemm/group_layernorm + multiple cat.
    We can have
    y1, y2, y3, y4 = group_layernorm([x1, x2, x3, x4])
    cat1 = concatenate(y1, y2)
    cat2 = concatenate(y3, y4)
    """
    sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
    for op in sorted_ops:
        if op._attrs["op"] != "concatenate":
            continue
        cat_op = op
        # check cat_dim is static
        cat_dim = cat_op._attrs["concat_dim"]
        if not transform_strided_ops_utils.cat_split_dim_is_static(cat_op, cat_dim):
            continue

        cat_inputs = cat_op._attrs["inputs"]
        cat_output = cat_op._attrs["outputs"][0]

        # We can have output_list = some_op()(x, y, z); y = concatenate()(output_list).
        # Depending on whether we create a copy or not in concatenate(),
        # the cat inputs and some_op outputs may point to the same list.
        # This copy adds extra safety to make sure that it doesn't happen.
        cat_op._attrs["inputs"] = cat_op._attrs["inputs"].copy()

        cat_inputs_to_remove = []
        for idx, cat_input in enumerate(cat_inputs):
            src_ops = list(cat_input.src_ops())
            if len(src_ops) != 1 or len(cat_input.dst_ops()) != 1:
                continue
            if cat_input._attrs["is_output"]:
                continue
            strided_op = src_ops[0]
            if not _is_supported_op(strided_op):
                continue

            # get output_idx into strided_op
            strided_idx = get_tensor_index(strided_op._attrs["outputs"], cat_input)

            if not _is_valid_for_fusion(strided_op, cat_op, strided_idx):
                continue

            # Update "outputs" of the strided_op, and "input_masks" of the cat op.
            # Note that we do not update "inputs" of the cat op here, to make sure
            # that we could find input indexes correctly for other tensors.
            if cat_op._attrs["inputs"].count(cat_input) > 1:
                # We do not support the case that a strided op write to multiple outputs.
                # TODO: Add multi-output later.
                continue

            offset = 0

            # cat's inputs may have been updated for cases like view_op + cat.
            # So, we need to retrieve original shapes from its input accessors.
            cat_input_accessors = cat_op._attrs["input_accessors"]
            # This pass must run before any other pass that remove cat inputs, like
            # _fuse_strided_op_reshape_cat
            for orig_i in range(idx):
                input_accessor = cat_input_accessors[orig_i]
                # TODO: Add dynamic shape support.
                offset += input_accessor.original_shapes[cat_dim].value()

            cat_inputs_to_remove.append(idx)

            strided_op._attrs["output_accessors"][strided_idx].update_base_tensor(
                cat_output, cat_dim, offset
            )

            cat_output._attrs["src_ops"].add(strided_op)

            output_tensor = strided_op._attrs["outputs"][strided_idx]
            strided_op._attrs["outputs"][strided_idx] = cat_output
            transform_utils.remove_tensor_from_sorted_graph(output_tensor)

            # Note that we have to update strided_op's epilogue_alignment
            # in the backend codegen where we modify the generate gemm
            # instance string with the updated epilogue_alignment.
            # The reason is similar to the one with slice + gemm.
            # For a problem size, the max alignment value returned from
            # output_accessor may vary for different concat cases.
            # If we update the strided_op's epilogue_alignment at this point,
            # the updated epilogue_alignment value will be cached in our
            # profiler database. Next time, when we get a different
            # updated epilogue_alignment value for the same problem size,
            # we would end up with not being able to find a matching
            # config or hit misalignment failures at runtime with invalid
            # alignment values.

        cat_op.remove_input_at(cat_inputs_to_remove)

    return transform_utils.sanitize_sorted_graph(sorted_graph)
Fuse strided_ops and cat op. One special case is group_gemm/group_layernorm + multiple cat.
We can have
y1, y2, y3, y4 = group_layernorm([x1, x2, x3, x4])
cat1 = concatenate(y1, y2)
cat2 = concatenate(y3, y4)
python/aitemplate/compiler/transform/transform_strided_ops.py