def _remove_make_jagged_ops(
    make_jagged_metadata: Dict[IntVar, List[MakeJaggedMetaData]],
    graph_inputs: Set[Tensor],
    graph_outputs: Set[Tensor],
):
    """Remove the make_jagged ops from the graph where possible.

    The individual make_jagged ops scattered over the graph are removed,
    to be further replaced by a single make_jagged instance, per total_length
    dimension, applied to all inputs with the total_length dimension at once.
    The ops are considered group by group, where group is formed from
    the ops with the same total_length dimension in the source Tensors.

    The make_jagged ops in the group are not removed (and the respective
    total_length key is popped from the make_jagged_metadata) if:

        1. There is only one make_jagged op in the group.

        2. There is a make_jagged op in the group connecting a
           graph input to a graph output: can't be eliminated.

        3. The total_length dimension representing the group is
           not present in any of the graph inputs' shape.

    In other cases, all make_jagged ops in the grpup are removed from the graph
    (and the respective total_length key is kept in the make_jagged_metadata).
    """
    for total_length in list(make_jagged_metadata.keys()):
        make_jagged_group = make_jagged_metadata[total_length]
        assert len({d.jagged_int_var for d in make_jagged_group}) == 1, (
            "All make_jagged ops applied to the sources with the "
            "same total_length must produce the same jagged_int_var."
        )  # this includes offsets identity check internally

        if len(make_jagged_group) == 1:
            _LOGGER.debug(
                "There is only one make_jagged op in the group "
                f"with {total_length=}: skipping the group."
            )
            make_jagged_metadata.pop(total_length)
            continue

        has_input_to_output_op = False
        for data in make_jagged_group:
            if any(s in graph_inputs for s in data.sources_list) and any(
                o in graph_outputs for o in data.outputs
            ):
                has_input_to_output_op = True
                break
        if has_input_to_output_op:
            _LOGGER.debug(
                "There is a make_jagged op in the group with "
                f"{total_length=} that maps a graph input to "
                "a graph output: skipping the group."
            )
            make_jagged_metadata.pop(total_length)
            continue

        graph_input_with_total_length = False
        for inp in graph_inputs:
            shape = inp._attrs["shape"]
            if shape and shape[0] == total_length:
                graph_input_with_total_length = True
                break
        if not graph_input_with_total_length:
            _LOGGER.debug(
                "None of the graph inputs has the first dimension "
                f"equal to {total_length=}: skipping the group."
            )
            make_jagged_metadata.pop(total_length)
            continue

        _LOGGER.debug(
            f"Removing {len(make_jagged_group)} make_jagged ops "
            f"in the group with {total_length=} from the graph."
        )
        for data in make_jagged_group:
            for source, output in zip(data.sources_list, data.outputs):
                replace_tensor(output, source)
                remove_dst_op_from_tensor(source, data.op)
Remove the make_jagged ops from the graph where possible.

    The individual make_jagged ops scattered over the graph are removed,
    to be further replaced by a single make_jagged instance, per total_length
    dimension, applied to all inputs with the total_length dimension at once.
    The ops are considered group by group, where group is formed from
    the ops with the same total_length dimension in the source Tensors.

    The make_jagged ops in the group are not removed (and the respective
    total_length key is popped from the make_jagged_metadata) if:

        1. There is only one make_jagged op in the group.

        2. There is a make_jagged op in the group connecting a
           graph input to a graph output: can't be eliminated.

        3. The total_length dimension representing the group is
           not present in any of the graph inputs' shape.

    In other cases, all make_jagged ops in the grpup are removed from the graph
    (and the respective total_length key is kept in the make_jagged_metadata).
    
python/aitemplate/compiler/transform/dedup_make_jagged_ops.py