def _fuse_split_and_strided_op(sorted_graph: List[Tensor]) -> List[Tensor]:
    """Fuse split and any op that supports strided inputs. This pass requires
    that all of the outputs can be fused into the next ops so that split op
    can be eliminated. Partial fusion is not supported yet.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        Input graph

    Returns
    -------
    List[Tensor]
        Fused graph
    """
    sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
    for op in sorted_ops:
        op_type = op._attrs["op"]
        if op_type != "split":
            continue
        split_op = op
        split_dim = split_op._attrs["split_dim"]

        split_input = split_op._attrs["inputs"][0]
        # split_dim must be static
        if not transform_strided_ops_utils.cat_split_dim_is_static(split_op, split_dim):
            continue

        outputs = split_op._attrs["outputs"]
        can_fuse_split = True

        stride = get_stride(split_input, split_dim)
        # offset on the split dim, which is different from the real offset
        dim_offset = 0
        output_offsets = []
        # We apply padding to bmm before this fuse_split pass. However, we may
        # still have mis-aligned accesses caused by offsets. This _check_alignment
        # filters out all bad cases.
        total_elems_from_split_dim = (
            stride * split_input._attrs["shape"][split_dim].value()
        )
        for output in outputs:
            can_fuse_split &= len(output.dst_ops()) > 0 and all(
                (
                    _is_supported_op(next_op._attrs["op"])
                    # need to pass the real offset to alignment checker
                    and _check_alignment(
                        next_op, dim_offset * stride, total_elems_from_split_dim
                    )
                    and len(output.dst_ops()) == 1
                )
                or (
                    next_op._attrs["op"] == "concatenate"
                    and next_op._attrs["concat_dim"] == split_dim
                    and _check_alignment(
                        next_op, dim_offset * stride, total_elems_from_split_dim
                    )
                )
                for next_op in output.dst_ops()
            )
            for next_op in output.dst_ops():
                for idx, input in enumerate(next_op._attrs["inputs"]):
                    if input == output:
                        can_fuse_split = can_fuse_split and (
                            transform_strided_ops_utils.gemm_stride_checker(
                                next_op._attrs["input_accessors"][idx], split_dim
                            )
                        )
            output_offsets.append(dim_offset)
            dim_offset += output._size(split_dim).value()

        if not can_fuse_split:
            continue

        _LOGGER.debug("Remove split from graph")
        split_input.dst_ops().remove(split_op)

        for output, offset in zip(outputs, output_offsets):
            for next_op in output.dst_ops():
                for idx, input in enumerate(next_op._attrs["inputs"]):
                    if input == output:
                        next_op._attrs["input_accessors"][idx].update_base_tensor(
                            split_input, split_dim, offset
                        )
                        # update the graph
                        next_op._attrs["inputs"][idx] = split_input
                split_input.dst_ops().add(next_op)

        # remove split op
        for output in outputs:
            transform_utils.remove_tensor_from_sorted_graph(output)
    return transform_utils.sanitize_sorted_graph(sorted_graph)
Fuse split and any op that supports strided inputs. This pass requires
    that all of the outputs can be fused into the next ops so that split op
    can be eliminated. Partial fusion is not supported yet.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        Input graph

    Returns
    -------
    List[Tensor]
        Fused graph
    
python/aitemplate/compiler/transform/fuse_split.py