[Pallas] Add block-sparse kernel tutorial

jax-ml · Sep 16, 2024 · 2e6df02 · 2e6df02
1 parent 29163fc
commit 2e6df02
Show file tree

Hide file tree

Showing 7 changed files with 1,306 additions and 0 deletions.
diff --git a/docs/_static/pallas/sparse/block_coo.svg b/docs/_static/pallas/sparse/block_coo.svg
diff --git a/docs/_static/pallas/sparse/prefetch_map.svg b/docs/_static/pallas/sparse/prefetch_map.svg
diff --git a/docs/_static/pallas/sparse/sparse_matmul.svg b/docs/_static/pallas/sparse/sparse_matmul.svg
diff --git a/docs/conf.py b/docs/conf.py
@@ -134,6 +134,7 @@ def _do_not_evaluate_in_jax(
     'pallas/quickstart.md',
     'pallas/tpu/pipelining.md',
     'pallas/tpu/distributed.md',
+    'pallas/tpu/sparse.md',
     'pallas/tpu/matmul.md',
     'jep/9407-type-promotion.md',
     'autodidax.md',
@@ -224,6 +225,7 @@ def _do_not_evaluate_in_jax(
     'pallas/quickstart.*',
     'pallas/tpu/pipelining.*',
     'pallas/tpu/distributed.*',
+    'pallas/tpu/sparse.*',
     'pallas/tpu/matmul.*',
     'sharded-computation.*',
     'distributed_data_loading.*'

diff --git a/docs/pallas/tpu/index.rst b/docs/pallas/tpu/index.rst
@@ -9,4 +9,6 @@ TPU specific documentation.
    details
    pipelining
    matmul
+   sparse
    distributed
+
diff --git a/docs/pallas/tpu/sparse.ipynb b/docs/pallas/tpu/sparse.ipynb
diff --git a/docs/pallas/tpu/sparse.md b/docs/pallas/tpu/sparse.md