mirror of
https://github.com/ROCm/jax.git
synced 2025-04-17 04:16:07 +00:00
46e65b5982
I've left the TPU backend docs a stub for now. Hopefully, someone working on Pallas TPU can fill them in later.
42 lines
589 B
ReStructuredText
42 lines
589 B
ReStructuredText
``jax.experimental.pallas.mosaic_gpu`` module
|
|
=============================================
|
|
|
|
.. automodule:: jax.experimental.pallas.mosaic_gpu
|
|
|
|
Classes
|
|
-------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
Barrier
|
|
GPUBlockSpec
|
|
GPUCompilerParams
|
|
GPUMemorySpace
|
|
TilingTransform
|
|
TransposeTransform
|
|
WGMMAAccumulatorRef
|
|
|
|
Functions
|
|
---------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
copy_gmem_to_smem
|
|
copy_smem_to_gmem
|
|
wait_barrier
|
|
wait_smem_to_gmem
|
|
wgmma
|
|
wgmma_wait
|
|
|
|
Aliases
|
|
-------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
ACC
|
|
GMEM
|
|
SMEM
|