xla

002f27a8 - [torchax] Fixes test for kthvalue pytorch/xla#7458 (#9223)