mirror of
https://github.com/ROCm/jax.git
synced 2025-04-18 12:56:07 +00:00
37 lines
428 B
ReStructuredText
37 lines
428 B
ReStructuredText
``jax.experimental.pallas`` module
|
|
==================================
|
|
|
|
.. automodule:: jax.experimental.pallas
|
|
|
|
Classes
|
|
-------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
BlockSpec
|
|
Slice
|
|
|
|
Functions
|
|
---------
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
pallas_call
|
|
program_id
|
|
num_programs
|
|
|
|
load
|
|
store
|
|
swap
|
|
|
|
atomic_and
|
|
atomic_add
|
|
atomic_cas
|
|
atomic_max
|
|
atomic_min
|
|
atomic_or
|
|
atomic_xchg
|
|
|
|
debug_print |