``jax.experimental.pallas`` module
==================================

.. automodule:: jax.experimental.pallas

Backends
--------

.. toctree::
    :maxdepth: 1

    Pallas TPU (TensorCore) <jax.experimental.pallas.tpu>
    Pallas MGPU <jax.experimental.pallas.mosaic_gpu>
    Triton <jax.experimental.pallas.triton>

Classes
-------

.. autosummary::
  :toctree: _autosummary

  BlockSpec
  GridSpec
  Slice

Functions
---------

.. autosummary::
  :toctree: _autosummary

  core_map
  kernel
  pallas_call
  program_id
  num_programs

  cdiv
  dslice
  empty
  empty_like

  broadcast_to
  debug_check
  debug_print
  dot
  get_global
  loop
  multiple_of
  run_scoped
  when
  with_scoped

Synchronization
---------------

.. autosummary::
  :toctree: _autosummary

  semaphore_read
  semaphore_signal
  semaphore_wait