jax

dd96fee0 - [Pallas/interpreter] Add an internal consistency check at the end of executing interpreted GPU kernels.

[Pallas/interpreter] Add an internal consistency check at the end of executing interpreted GPU kernels. All barriers should have been deallocated again when the interpreted kernel finishes executing. This is an invariant that should be maintained by the interpreter internally. Hence we `assert` on this invariant. Note: By ensuring that all barriers have been deallocated, we guarantee that we check that threads observe all barrier completions (if any) since this check is performed during barrier deallocation. PiperOrigin-RevId: 862967671