xla

8b8c1a00 - [torchax] Fixes test for cat pytorch/xla#7398 (#9047)