jax.experimental.pallas.mosaic_gpu.wgmma_wait

jax.experimental.pallas.mosaic_gpu.wgmma_wait(n)

等待直到不超过 n 个正在进行的 WGMMA 操作。

参数:

n (int)