jax.experimental.pallas.mosaic_gpu.wgmma_wait

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

等待直到在飞行中 WGMMA 操作不超过 n 个。

参数:

n (int)