def dedup_make_jagged_ops(
    sorted_graph: List[Tensor],
    workdir: str = None,
) -> List[Tensor]:
    """Deduplicate make_jagged ops in the graph.

    The rationale is to eliminate redundant offset validation as
    well as make the implicit jagged Tensors (sources) in the graph
    explicit, by replacing their total_length dimension with the
    corresponding JaggedIntVar.

    The pass is performed in the following steps:

        1. Collect the metadata of the existing make_jagged ops.
        2. Remove make_jagged ops from the graph where possible.
        3. Apply new make_jagged ops to the (bundled) source inputs.
        4. Replace total_length dimensions with new JaggedIntVars.

    See the docstrings of the individual steps' helper functions
    above for more details.
    """
    make_jagged_metadata = _get_make_jagged_metadata(sorted_graph)

    if not make_jagged_metadata:
        _LOGGER.debug("No make_jagged ops in the graph: skipping.")
        return sorted_graph

    graph_inputs = {t for t in sorted_graph if t._attrs["is_input"]}
    graph_outputs = {t for t in sorted_graph if t._attrs["is_output"]}

    _remove_make_jagged_ops(
        make_jagged_metadata,
        graph_inputs,
        graph_outputs,
    )

    if not make_jagged_metadata:
        _LOGGER.debug(
            "There are make_jagged ops in the graph, "
            "but nothing to deduplicate: skipping."
        )
        return sorted_graph

    # drop the removed make_jagged outputs
    sorted_graph = sanitize_sorted_graph(sorted_graph)

    new_jagged_int_vars = _apply_make_jagged_to_inputs(
        make_jagged_metadata,
        sorted_graph,
        graph_inputs,
    )
    _replace_total_length_with_jagged_int_var(
        new_jagged_int_vars,
        sorted_graph,
        graph_inputs,
    )

    # sort the new make_jagged outputs
    sorted_graph = toposort(sorted_graph)
    # name the new tensors + do sanity check
    sorted_graph = sanitize_sorted_graph(sorted_graph)

    return sorted_graph
Deduplicate make_jagged ops in the graph.

The rationale is to eliminate redundant offset validation as
well as make the implicit jagged Tensors (sources) in the graph
explicit, by replacing their total_length dimension with the
corresponding JaggedIntVar.

The pass is performed in the following steps:

1. Collect the metadata of the existing make_jagged ops.
2. Remove make_jagged ops from the graph where possible.
3. Apply new make_jagged ops to the (bundled) source inputs.
4. Replace total_length dimensions with new JaggedIntVars.

See the docstrings of the individual steps' helper functions
above for more details.
python/aitemplate/compiler/transform/dedup_make_jagged_ops.py