diff --git a/Tensor Puzzlers.ipynb b/Tensor Puzzlers.ipynb
index 52b7078..708646d 100644
--- a/Tensor Puzzlers.ipynb	
+++ b/Tensor Puzzlers.ipynb	
@@ -2283,14 +2283,20 @@
     "def repeat(a: TT[\"i\"], d: TT[1]) -> TT[\"d\", \"i\"]:\n",
     "    raise NotImplementedError\n",
     "\n",
-    "test_repeat = make_test(\"repeat\", repeat, repeat_spec, constraint=constraint_set)\n",
-    "\n",
-    "\n",
-    "# ## Puzzle 21 - bucketize\n",
-    "#\n",
-    "# Compute [bucketize](https://pytorch.org/docs/stable/generated/torch.bucketize.html)"
+    "test_repeat = make_test(\"repeat\", repeat, repeat_spec, constraint=constraint_set)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Puzzle 21 - bucketize\n",
+    "\n",
+    "Compute [bucketize](https://pytorch.org/docs/stable/generated/torch.bucketize.html)"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
   {
    "cell_type": "code",
    "execution_count": 45,
@@ -2349,15 +2355,20 @@
     "    raise NotImplementedError\n",
     "\n",
     "test_bucketize = make_test(\"bucketize\", bucketize, bucketize_spec,\n",
-    "                           constraint=constraint_set)\n",
-    "\n",
-    "\n",
-    "#\n",
-    "# # Speed Run Mode!\n",
-    "#\n",
-    "# What is the smallest you can make each of these?"
+    "                           constraint=constraint_set)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Speed Run Mode!\n",
+    "\n",
+    "What is the smallest you can make each of these?"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
   {
    "cell_type": "code",
    "execution_count": 46,
@@ -2397,7 +2408,8 @@
    "source": [
     "import inspect\n",
     "fns = (ones, sum, outer, diag, eye, triu, cumsum, diff, vstack, roll, flip,\n",
-    "       compress, pad_to, sequence_mask, bincount, scatter_add)\n",
+    "       compress, pad_to, sequence_mask, bincount, scatter_add, flatten,\n",
+    "       linspace, heaviside, repeat, bucketize)\n",
     "\n",
     "for fn in fns:\n",
     "    lines = [l for l in inspect.getsource(fn).split(\"\\n\") if not l.strip().startswith(\"#\")]\n",