def _apply_make_jagged_to_inputs(
    make_jagged_metadata: Dict[IntVar, List[MakeJaggedMetaData]],
    sorted_graph: List[Tensor],
    graph_inputs: Set[Tensor],
) -> Dict[IntVar, JaggedIntVar]:
    """Apply new make_jagged ops to the (bundled) input source Tensors.

    For each group of make_jagged ops that removed from the graph,
    a new make_jagged op is applied to all graph inputs with the
    corresponding total_length dimension. This way, the source Tensors
    are converted to jagged Tensors right from the "beginning" of the
    graph and can be used as jagged Tensors downstream.

    Two points are worth mentioning:

        1. Due to the fact that the new make_jagged op is applied to
           *all* source inputs with the total_length dimension, it is
           guaranteed that the offsets validation performed by the
           make_jagged op's back-end will run before any of the
           resulting jagged Tensors can be used downstream.

        2. Because a single make_jagged op is applied to multiple
           graph inputs, the make_jagged op's back-end kernel will
           be launched only once to validate the offsets (the latter
           are the same for every source input). This optimizes out
           redundant validation of the same offsets.

    The mapping of each total_length to the new JaggedIntVar (produced
    by the corresponding new make_jagged op) is returned.
    """
    new_jagged_int_vars = {}
    for total_length, make_jagged_group in make_jagged_metadata.items():
        sources_list = []
        for inp in graph_inputs:
            shape = inp._attrs["shape"]
            if shape and shape[0] == total_length:
                sources_list.append(inp)

        _LOGGER.debug(
            "Adding a single make_jagged op for the source inputs "
            f"{[source._attrs['name'] for source in sources_list]}."
        )

        data = make_jagged_group[0]
        new_make_jagged_op = make_jagged(
            batch_dim=data.jagged_int_var.batch_dim(),
            jagged_dims=data.jagged_int_var.jagged_dims(),
            check_sequence_lengths=all(
                d.op._attrs["check_sequence_lengths"] for d in make_jagged_group
            ),
        )
        jagged_tensors = new_make_jagged_op(
            source=sources_list,
            offsets_list=data.offsets_list,
        )
        jagged_int_var = jagged_tensors[0]._attrs["shape"][0]
        new_jagged_int_vars[total_length] = jagged_int_var

        for source, jagged in zip(sources_list, jagged_tensors):
            for op in source._attrs["dst_ops"]:
                if op is not new_make_jagged_op:
                    replace_tensor_for_op(op, source, jagged)

        sorted_graph.extend(jagged_tensors)

    return new_jagged_int_vars
Apply new make_jagged ops to the (bundled) input source Tensors.

For each group of make_jagged ops that removed from the graph,
a new make_jagged op is applied to all graph inputs with the
corresponding total_length dimension. This way, the source Tensors
are converted to jagged Tensors right from the "beginning" of the
graph and can be used as jagged Tensors downstream.

Two points are worth mentioning:

1. Due to the fact that the new make_jagged op is applied to
*all* source inputs with the total_length dimension, it is
guaranteed that the offsets validation performed by the
make_jagged op's back-end will run before any of the
resulting jagged Tensors can be used downstream.

2. Because a single make_jagged op is applied to multiple
graph inputs, the make_jagged op's back-end kernel will
be launched only once to validate the offsets (the latter
are the same for every source input). This optimizes out
redundant validation of the same offsets.

The mapping of each total_length to the new JaggedIntVar (produced
by the corresponding new make_jagged op) is returned.
python/aitemplate/compiler/transform/dedup_make_jagged_ops.py