def _fix_op_inputs_outputs(
    subgraph: List[Tensor], name_to_new_tensor: Dict[str, Tensor]
) -> None:
    """
    This is an unfortunate hack made necessary by the following:

    1) When constructing the constant folding subgraph, the most understandable
       thing to do is create *new* tensors so we can modify their attributes without
       affecting the original graph.
    2) However, the inputs of each tensor's src and dst ops need to be wired up to
       the new tensors since the memory planning pass will traverse the graph through those attributes.

    So, we store the mapping from tensor name to its corresponding subgraph tensor and the tensor in
    original graph.

    Before we do memory planning for constant folding, we call:
      _fix_op_inputs_outputs(subgraph, name_to_constant_folding_tensor)

    And then afterwards we restore everything with:
      _fix_op_inputs_outputs(subgraph, name_to_original_tensor)

    It would be nice if we could deep copy the src and dst ops when we create new tensors so we can
    skip the restoration step. But this is not implemented and not trivial. Thankfully, this function
    is not too hard to understand once the rationale behind it is understood.
    """
    ops = graph_utils.get_sorted_ops(subgraph)
    for op in ops:
        op._attrs["inputs"] = [
            name_to_new_tensor[tensor._attrs["name"]] for tensor in op._attrs["inputs"]
        ]

        op._attrs["outputs"] = [
            name_to_new_tensor[tensor._attrs["name"]] for tensor in op._attrs["outputs"]
        ]
This is an unfortunate hack made necessary by the following:

1) When constructing the constant folding subgraph, the most understandable
thing to do is create *new* tensors so we can modify their attributes without
affecting the original graph.
2) However, the inputs of each tensor's src and dst ops need to be wired up to
the new tensors since the memory planning pass will traverse the graph through those attributes.

So, we store the mapping from tensor name to its corresponding subgraph tensor and the tensor in
original graph.

Before we do memory planning for constant folding, we call:
_fix_op_inputs_outputs(subgraph, name_to_constant_folding_tensor)

And then afterwards we restore everything with:
_fix_op_inputs_outputs(subgraph, name_to_original_tensor)

It would be nice if we could deep copy the src and dst ops when we create new tensors so we can
skip the restoration step. But this is not implemented and not trivial. Thankfully, this function
is not too hard to understand once the rationale behind it is understood.
python/aitemplate/compiler/transform/constant_folding.py