torch.cuda.synchronize counterpart / Using cupy.cuda.Device.synchronize

[jaywonchung]

Hi,

I'm trying to see if I can create a generic DL-framework-agnostic way to synchronize CPU and GPU executions (for use in a framework I'm writing).
In PyTorch, users can call torch.cuda.synchronize() pretty much anywhere for synchronization.
Would there be a counterpart in JAX?

I am aware of id_tap and block_until_ready, but they require a reference to the result of computation, unlike torch.cuda.synchronize().

If JAX itself doesn't have such a primitize, what would happen if I instantiate a cupy.cuda.Device and then call synchronize on it? Would this somehow disrupt what JAX is doing?

Thanks a lot.