Check if the current Pallas kernels have enough test coverage.

[vanbasten23]

@miladm

This may boil down to test on various TPU hardwares, fuzzy tests, etc.

[miladm]

cc @ysiraichi as he is working on adding more coverage tests to the torchxla stack.

Hopefully the two of you can work together on this.

This kind of test is useful for all of our tests (beyond kernels).

Let's start with Pallas kernel and expand after @ysiraichi.