def replace_tensor_for_op(target_op: Operator, old_tensor: Tensor, new_tensor: Tensor):
    """
    Replaces all references to the old_tensor with the new_tensor in target_op.
    """

    dst_ops = list(old_tensor._attrs["dst_ops"])
    for op in dst_ops:
        if op is target_op:
            op.replace_input_tensor(old_tensor, new_tensor)
            new_tensor._attrs["dst_ops"].add(op)
            old_tensor._attrs["dst_ops"].remove(op)
            break

    Replaces all references to the old_tensor with the new_tensor in target_op.
    
python/aitemplate/compiler/transform/transform_utils.py