jax

e6299e2e - Expose optimization levels to Pallas

Expose optimization levels to Pallas This flag is only used for ``SC_*_SUBCORE`` kernels and it implicitly defaults to O3. Going below O2 will currently result in compilation failures. PiperOrigin-RevId: 940557264