xla

87df88c6 - Workaround all_gather graph breaks