def constant_folding(
    sorted_graph: List[Tensor],
    workdir: str,
    model_name: str,
) -> Tuple[List[Tensor], List[Tuple[str, str]], List[Tensor]]:
    """
    Fold and propagate constants.

    This pass looks for ops that have inputs which can be determined
    at compile time. It evaluates them, then puts the new constants
    back into the graph with bound data. The old ops are eliminated.

    This pass actually compiles and runs an AIT runtime. If there are
    any problems (e.g. due to buggy ops), the constant folding is
    aborted and the graph is returned unchanged. All generated code
    is stored in workdir/constant_folding.
    """
    new_constants, file_pairs, constant_folding_inputs = _constant_folding_impl(
        sorted_graph, workdir, model_name
    )

    # Replace ops with their folded values.
    for idx, tensor in enumerate(sorted_graph):
        name = tensor._attrs["name"]
        if name in new_constants:
            new_tensor = new_constants[name]
            replace_tensor(tensor, new_tensor)
            sorted_graph[idx] = new_tensor

    # Eliminate constants that are no longer used
    compiler.transform.remove_unused_ops(sorted_graph)
    return (
        compiler.transform.transform_utils.sanitize_sorted_graph(sorted_graph),
        file_pairs,
        constant_folding_inputs,
    )
Fold and propagate constants.

This pass looks for ops that have inputs which can be determined
at compile time. It evaluates them, then puts the new constants
back into the graph with bound data. The old ops are eliminated.

This pass actually compiles and runs an AIT runtime. If there are
any problems (e.g. due to buggy ops), the constant folding is
aborted and the graph is returned unchanged. All generated code
is stored in workdir/constant_folding.
python/aitemplate/compiler/transform/constant_folding.py