jax

d4d63098 - [ROCm][Pallas] Skip fused attention fwd tests that exceed device resources.

[ROCm][Pallas] Skip fused attention fwd tests that exceed device resources. Similar to the bwd fused attention case, instead of keeping a hardcoded list of parameter combinations, we catch RESOURCE_EXHAUSTED errors from the compiler and skip the test dynamically. This handles cases where a kernel's shared memory requirements exceed the hardware's limits. (cherry picked from commit e76d682d4f61ce3ef04f9a93623ea433666d312b)