Fix handling of tuples in CreateCopyWithNewLayout.
If the layout of a single tensor in a tuple is different from its use, then CreateCopyWithNewLayout will do a deep copy of the entire tuple. Not only does this operation create unnecessary copies of elements where the layout is the same, it will throw an error if the tuple contains elements like token[] that cannot be copied. As a result, layout assignment on TPU occassionally causes mysterious compilation failures for code that runs correctly on CPU and GPU. PiperOrigin-RevId: 215615731
Loading
Please sign in to comment