jax

38d52a19 - [mosaic_gpu] Force flush all cupti activity, then unsubscribe.

[mosaic_gpu] Force flush all cupti activity, then unsubscribe. With default flushing, it is possible for events to be missed. We should only unsubscribe after we are finished with cupti. PiperOrigin-RevId: 737939327