* Don't copy on-host and on-device shapes locally.
* Use ForEachMutableElement rather than the iterators, as it is much quicker. There is still room for improvement; ForEachMutableElement is linear in the number of nodes in the shape tree but we want to be linear in the number of nodes in the sub shape tree. But I feel this is a good enough improvement. PiperOrigin-RevId: 195384423
Loading
Please sign in to comment