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

.. automodule:: jax.experimental.pallas.tpu

Classes
-------

.. autosummary::
   :toctree: _autosummary

   ChipVersion
   CompilerParams
   GridDimensionSemantics
   MemorySpace
   PrefetchScalarGridSpec
   SemaphoreType
   TpuInfo

Functions
---------

.. autosummary::
   :toctree: _autosummary

   load
   store

Communication
-------------

.. autosummary::
   :toctree: _autosummary

   async_copy
   async_remote_copy
   make_async_copy
   make_async_remote_copy
   sync_copy

Pipelining
----------

.. autosummary::
   :toctree: _autosummary

   BufferedRef
   BufferedRefBase
   emit_pipeline
   emit_pipeline_with_allocations

Pseudorandom Number Generation
------------------------------

.. autosummary::
   :toctree: _autosummary

   prng_seed
   sample_block
   stateful_bernoulli
   stateful_bits
   stateful_normal
   stateful_uniform
   to_pallas_key

Interpret Mode
--------------

.. autosummary::
   :toctree: _autosummary

   force_tpu_interpret_mode
   InterpretParams
   reset_tpu_interpret_mode_state
   set_tpu_interpret_mode

Miscellaneous
-------------

.. autosummary::
   :toctree: _autosummary

   core_barrier
   get_barrier_semaphore
   get_tpu_info
   is_tpu_device
   run_on_first_core
   with_memory_space_constraint
