[Pallas] PoC Integration (#6340)
Summary:
This is PoC for Pallas integration. Currently, it can run Pallas kernels that take arbitrary tensors as inputs and output a single tensor. The design doc is here: go/pytorch-xla-pallas.
Test Plan:
PJRT_DEVICE=TPU python test/test_operations.py -v -k test_tpu_custom_call
