From 91f16419bb3b1670cf31c55a4e8dde194bb73f53 Mon Sep 17 00:00:00 2001
From: Dongseong Hwang <dongseong@google.com>
Date: Mon, 23 Sep 2024 09:06:35 -0700
Subject: [PATCH] Fix errata in block-sparse kernel tutorial.

Correct M//blk_M to N//blk_N. It was ok because both values happen to be same.
In addition, grid order is (num_blocks, j) as 'num_blocks' replaces 'i'.

PiperOrigin-RevId: 677817478
---
 docs/pallas/tpu/sparse.ipynb | 8 ++++----
 docs/pallas/tpu/sparse.md    | 8 ++++----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/docs/pallas/tpu/sparse.ipynb b/docs/pallas/tpu/sparse.ipynb
index 909103273..6666b9e3e 100644
--- a/docs/pallas/tpu/sparse.ipynb
+++ b/docs/pallas/tpu/sparse.ipynb
@@ -312,13 +312,13 @@
     "    o_ref[...] = accum_scratch[...].astype(o_ref.dtype)\n",
     "\n",
     "\n",
-    "def x_map(j, blk_idx, blk_idxs_i, blk_idxs_k):\n",
+    "def x_map(blk_idx, j, blk_idxs_i, blk_idxs_k):\n",
     "  del j, blk_idxs_i, blk_idxs_k\n",
     "  return (blk_idx, 0, 0)\n",
-    "def y_map(j, blk_idx, blk_idxs_i, blk_idxs_k):\n",
+    "def y_map(blk_idx, j, blk_idxs_i, blk_idxs_k):\n",
     "  del blk_idxs_i\n",
     "  return (blk_idxs_k[blk_idx], j)\n",
-    "def o_map(j, blk_idx, blk_idxs_i, blk_idxs_k):\n",
+    "def o_map(blk_idx, j, blk_idxs_i, blk_idxs_k):\n",
     "  del blk_idxs_k\n",
     "  return (blk_idxs_i[blk_idx], j)\n",
     "\n",
@@ -333,7 +333,7 @@
     "    num_scalar_prefetch=2,\n",
     "    # Note that while num_blocks is static here, Pallas does support\n",
     "    # dynamic grid sizes.\n",
-    "    grid=(M // blk_M, num_blocks),\n",
+    "    grid=(num_blocks, N // blk_N),\n",
     "    in_specs=[pl.BlockSpec((1, blk_M, blk_K), x_map),\n",
     "              pl.BlockSpec((blk_K, blk_N), y_map),\n",
     "              # Placeholder for a zeros-array used by input_output_aliases.\n",
diff --git a/docs/pallas/tpu/sparse.md b/docs/pallas/tpu/sparse.md
index 23e14bb9b..3bc662895 100644
--- a/docs/pallas/tpu/sparse.md
+++ b/docs/pallas/tpu/sparse.md
@@ -252,13 +252,13 @@ def dsd_kernel(idxs_i_ref, idxs_k_ref, # Scalar prefetch inputs.
     o_ref[...] = accum_scratch[...].astype(o_ref.dtype)
 
 
-def x_map(j, blk_idx, blk_idxs_i, blk_idxs_k):
+def x_map(blk_idx, j, blk_idxs_i, blk_idxs_k):
   del j, blk_idxs_i, blk_idxs_k
   return (blk_idx, 0, 0)
-def y_map(j, blk_idx, blk_idxs_i, blk_idxs_k):
+def y_map(blk_idx, j, blk_idxs_i, blk_idxs_k):
   del blk_idxs_i
   return (blk_idxs_k[blk_idx], j)
-def o_map(j, blk_idx, blk_idxs_i, blk_idxs_k):
+def o_map(blk_idx, j, blk_idxs_i, blk_idxs_k):
   del blk_idxs_k
   return (blk_idxs_i[blk_idx], j)
 
@@ -273,7 +273,7 @@ grid_spec = pltpu.PrefetchScalarGridSpec(
     num_scalar_prefetch=2,
     # Note that while num_blocks is static here, Pallas does support
     # dynamic grid sizes.
-    grid=(M // blk_M, num_blocks),
+    grid=(num_blocks, N // blk_N),
     in_specs=[pl.BlockSpec((1, blk_M, blk_K), x_map),
               pl.BlockSpec((blk_K, blk_N), y_map),
               # Placeholder for a zeros-array used by input_output_aliases.