def replace_tensor(old_tensor, new_tensor):
    """
    Replaces all references to the old_tensor with the new_tensor.
    """
    if old_tensor._attrs["is_output"]:
        new_tensor._attrs["is_output"] = True
        new_tensor._attrs["name"] = old_tensor._attrs["name"]
    dst_ops = list(old_tensor._attrs["dst_ops"])
    for op in dst_ops:
        op.replace_input_tensor(old_tensor, new_tensor)
        new_tensor._attrs["dst_ops"].add(op)
        old_tensor._attrs["dst_ops"].remove(op)
    remove_tensor_from_sorted_graph(old_tensor)
Replaces all references to the old_tensor with the new_tensor.
python/aitemplate/compiler/transform/transform_utils.py