jax

68972de0 - [pallas] Add lowering errors for block shapes that are not supported.

[pallas] Add lowering errors for block shapes that are not supported. Previously these errors came from Mosaic with less useful stack traces, and in the case of GPU we were getting a crash instead of an exception. PiperOrigin-RevId: 657184114