def _replace_total_length_with_jagged_int_var(
    new_jagged_int_vars: Dict[IntVar, JaggedIntVar],
    sorted_graph: List[Tensor],
    graph_inputs: Set[Tensor],
):
    """Replace total_length dimensions by the new JaggedIntVars.

    As we've removed the internal make_jagged ops from the graph and
    replaced their output jagged Tensors by the input source Tensors,
    the latter have lost their JaggedIntVars. Here we replace the
    total_length dimension in *every* non-input Tensor in the graph
    by the corresponding new JaggedIntVar (produced by the new
    make_jagged op applied to the bundled source inputs). This includes,
    but is not limited to, the source inputs of the make_jagged ops
    removed from within the graph in the beginning of the pass.
    """
    for total_length, new_jagged_int_var in new_jagged_int_vars.items():
        for tensor in sorted_graph:
            if tensor not in graph_inputs:
                shape = tensor._attrs["shape"]
                if shape and shape[0] == total_length:
                    shape[0] = new_jagged_int_var
Replace total_length dimensions by the new JaggedIntVars.

As we've removed the internal make_jagged ops from the graph and
replaced their output jagged Tensors by the input source Tensors,
the latter have lost their JaggedIntVars. Here we replace the
total_length dimension in *every* non-input Tensor in the graph
by the corresponding new JaggedIntVar (produced by the new
make_jagged op applied to the bundled source inputs). This includes,
but is not limited to, the source inputs of the make_jagged ops
removed from within the graph in the beginning of the pass.
python/aitemplate/compiler/transform/dedup_make_jagged_ops.py