def _extract_foldable_subgraph(
    sorted_graph: List[Tensor],
) -> Tuple[List[Tensor], Dict[str, Tensor], List[Tensor]]:
    """
    Extract a list of foldable nodes. A node is foldable if:
    * It has bound data, or
    * All of its inputs are foldable.

    The subgraph returned is just a list of Tensors. All foldable
    tensors that do not have bound data are marked as outputs in
    the subgraph. The original graph is not modified.

    All tensors that do not have bound data are marked as outputs.
    This is because we want to execute the subgraph and get all
    of the new constants. Only the ones that are actually needed are put
    back into the final graph.
    """
    foldable_node_names = set()
    foldable_ops = set()
    subgraph = []

    for tensor in sorted_graph:
        if tensor._attrs["is_input"] or tensor._attrs["skip_constant_folding"]:
            continue

        name = tensor._attrs["name"]
        if tensor._attrs["data"] is not None or tensor._attrs["is_param"]:
            foldable_node_names.add(name)
            subgraph.append(tensor)
            continue
        elif isinstance(tensor, IntVarTensor):
            continue
        foldable = all(
            inp._attrs["name"] in foldable_node_names
            for op in tensor._attrs["src_ops"]
            for inp in op._attrs["inputs"]
        )

        if foldable:
            foldable_node_names.add(name)
            subgraph.append(tensor)
            for op in tensor._attrs["src_ops"]:
                foldable_ops.add(op)

    def _is_used_by_non_foldable_op(tensor: Tensor) -> bool:
        for op in tensor._attrs["dst_ops"]:
            if op not in foldable_ops:
                return True
        return False

    def _is_used_by_foldable_op(tensor: Tensor) -> bool:
        for op in tensor._attrs["dst_ops"]:
            if op in foldable_ops:
                return True
        return False

    # Now figure out which tensors can be marked as outputs.
    filtered_subgraph = []
    name_to_new_tensor = {}
    name_to_old_tensor = {}
    constant_folding_inputs = []

    for tensor in subgraph:
        name = tensor._attrs["name"]
        new_tensor = None

        if not tensor._attrs["is_param"] and (
            _is_used_by_non_foldable_op(tensor) or tensor._attrs["is_output"]
        ):
            # Tensor is required outside of the subgraph, make it an output.
            # Parameters don't need to be marked as outputs in the
            # subgraph, we already know their values.
            new_tensor = _output_from_tensor(tensor)

        elif _is_used_by_foldable_op(tensor):
            # No need to append constants that are not used by any foldable ops.
            new_tensor = _non_output_from_tensor(tensor)
            if new_tensor._attrs["is_param"]:
                constant_folding_inputs.append(new_tensor)

        if new_tensor is not None:
            name_to_new_tensor[name] = new_tensor
            name_to_old_tensor[name] = tensor
            filtered_subgraph.append(new_tensor)

    _fix_op_inputs_outputs(filtered_subgraph, name_to_new_tensor)
    return filtered_subgraph, name_to_old_tensor, constant_folding_inputs
Extract a list of foldable nodes. A node is foldable if:
* It has bound data, or
* All of its inputs are foldable.

The subgraph returned is just a list of Tensors. All foldable
tensors that do not have bound data are marked as outputs in
the subgraph. The original graph is not modified.

All tensors that do not have bound data are marked as outputs.
This is because we want to execute the subgraph and get all
of the new constants. Only the ones that are actually needed are put
back into the final graph.
python/aitemplate/compiler/transform/constant_folding.py