jax.experimental.pallas.tpu.async_copy

jax.experimental.pallas.tpu.async_copy(src_ref, dst_ref, sem, *, priority=0, add=False)

Issues a DMA copying from src_ref to dst_ref.

Parameters:

• priority (int)

• add (bool)

Return type:

AsyncCopyDescriptor