Improve the implementation of HloSharding::GetSubSharding
The new implementation iterates the ShapeIndex and the flat HloSharding directly instead of constructing a ShapeTree what provides a significant performance benefit for large tuple shapes. PiperOrigin-RevId: 210705719
Loading
Please sign in to comment