jax

99a8b92a - Fixed Pallas Mosaic GPU tests

Fixed Pallas Mosaic GPU tests * Migrated to the new barrier APIs * Fixed scratch view casting logic, it previously didn't work for >1 view PiperOrigin-RevId: 655937541