xla

[Pallas] Make repeat_with_fixed_output_size not OOM on VMEM

#7145

Merged