.. _pallas: Pallas: a JAX kernel language ============================= Pallas is an extension to JAX that enables writing custom kernels for GPU and TPU. This section contains tutorials, guides and examples for using Pallas. See also the :class:`jax.experimental.pallas` module API documentation. .. warning:: Pallas is experimental and is changing frequently. See the :ref:`pallas-changelog` for the recent changes. .. toctree:: :caption: Guides :maxdepth: 2 quickstart design grid_blockspec .. toctree:: :caption: Platform Features :maxdepth: 2 tpu/index .. toctree:: :maxdepth: 1 CHANGELOG