xla

ca914452 - Add support for callable in torchax.interop.JittableModule.functional_call in the first parameter (#9451)