xla

0aac10e2 - [torchax] Fix functional.max_pool (#8814)