Merge pull request #546 from Routhleck/master

[docs] Optimize tutorial code in `operator_custom_with_taichi.ipynb` of documentations
brainpy · Nov 16, 2023 · 20c1f40 · 20c1f40
2 parents 5432a2b + a3f73df
commit 20c1f40
Showing 1 changed file with 12 additions and 12 deletions.
diff --git a/docs/tutorial_advanced/operator_custom_with_taichi.ipynb b/docs/tutorial_advanced/operator_custom_with_taichi.ipynb
@@ -4,7 +4,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Operator Customization with Numba"
+    "# Operator Customization with Taichi"
    ]
   },
   {
@@ -18,7 +18,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Brain dynamics is sparse and event-driven, however, proprietary operators for brain dynamics are not well abstracted and summarized. As a result, we are often faced with the need to customize operators. In this tutorial, we will explore how to customize brain dynamics operators using ti.\n",
+    "Brain dynamics is sparse and event-driven, however, proprietary operators for brain dynamics are not well abstracted and summarized. As a result, we are often faced with the need to customize operators. In this tutorial, we will explore how to customize brain dynamics operators using taichi.\n",
     "\n",
     "Start by importing the relevant Python package."
    ]
@@ -242,11 +242,11 @@
     "                  vector: ti.types.ndarray(ndim=1), \n",
     "                  weight: ti.types.ndarray(ndim=1), \n",
     "                  out: ti.types.ndarray(ndim=1)):\n",
-    "   weight_0 = weight[0]\n",
-    "   ti.loop_config(block_dim=64)\n",
-    "   for i, j in ti.ndrange(indices.shape[0], indices.shape[1]):\n",
-    "        if vector[i]:\n",
-    "            out[indices[i, j]] += weight_0\n",
+    "  weight_0 = weight[0]\n",
+    "  ti.loop_config(block_dim=64)\n",
+    "  for ij in ti.grouped(indices):\n",
+    "      if vector[ij[0]]:\n",
+    "          out[ij[1]] += weight_0\n",
     "\n",
     "prim = bm.XLACustomOp(cpu_kernel=event_ell_cpu, gpu_kernel=event_ell_gpu)\n",
     "\n",
@@ -503,11 +503,11 @@
     "                  vector: ti.types.ndarray(ndim=1), \n",
     "                  weight: ti.types.ndarray(ndim=1), \n",
     "                  out: ti.types.ndarray(ndim=1)):\n",
-    "   weight_0 = weight[0]\n",
-    "   ti.loop_config(block_dim=64)\n",
-    "   for i, j in ti.ndrange(indices.shape[0], indices.shape[1]):\n",
-    "        if vector[i]:\n",
-    "            out[indices[i, j]] += weight_0\n",
+    "  weight_0 = weight[0]\n",
+    "  ti.loop_config(block_dim=64)\n",
+    "  for ij in ti.grouped(indices):\n",
+    "      if vector[ij[0]]:\n",
+    "          out[ij[1]] += weight_0\n",
     "\n",
     "prim = bm.XLACustomOp(cpu_kernel=event_ell_cpu, gpu_kernel=event_ell_gpu)\n",
     "\n",