jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive#

jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive(barrier, collective_axis=None)[source]#

Tracks completion of a preceding tcgen05_mma call.

Parameters:

  • barrier (_Ref) – Barrier Ref for synchronizing with the tensor core. Must have orders_tensor_core set to True.

  • collective_axis (str | None) – The name of the cluster axis along which the MMA was performed if it was collective. The cluster axis should have a size of exactly 2, and must be on the minormost cluster axis.

See also

jax.experimental.pallas.mosaic_gpu.tcgen05_mma()