def _get_make_jagged_metadata(
    sorted_graph: List[Tensor],
) -> Dict[IntVar, List[MakeJaggedMetaData]]:
    """Collect metadata about the existing make_jagged ops in the graph.

    The MakeJaggedMetaData instances, one per make_jagged op, are grouped
    by the total_length dimension in the source input Tensors of the ops.
    In case of multiple inputs, total_length dimension is the same in
    every input. The metadata is used further to inform the transformation.
    """
    metadata = {}
    for op in get_sorted_ops(sorted_graph):
        if op._attrs["op"] == "make_jagged":
            outputs = op._attrs["outputs"]
            jagged_int_var = outputs[0]._attrs["shape"][0]
            total_length = jagged_int_var.total_length()
            num_sources = op._attrs["num_sources"]
            if total_length not in metadata:
                metadata[total_length] = []
            metadata[total_length].append(
                MakeJaggedMetaData(
                    op=op,
                    sources_list=op._attrs["inputs"][:num_sources],
                    offsets_list=op._attrs["inputs"][num_sources:],
                    outputs=outputs,
                    jagged_int_var=jagged_int_var,
                )
            )

    return metadata
Collect metadata about the existing make_jagged ops in the graph.

The MakeJaggedMetaData instances, one per make_jagged op, are grouped
by the total_length dimension in the source input Tensors of the ops.
In case of multiple inputs, total_length dimension is the same in
every input. The metadata is used further to inform the transformation.
python/aitemplate/compiler/transform/dedup_make_jagged_ops.py