diff --git a/docs/2-einops-for-deep-learning.ipynb b/docs/2-einops-for-deep-learning.ipynb
index 8c89a887..4de58085 100644
--- a/docs/2-einops-for-deep-learning.ipynb
+++ b/docs/2-einops-for-deep-learning.ipynb
@@ -1161,6 +1161,8 @@
     "    Linear(16*5*5, 120), \n",
     "    ReLU(),\n",
     "    Linear(120, 10), \n",
+    "    # In flax, the {'axis': value} syntax for specifying values for axes is mandatory:\n",
+    "    # Rearrange('(b1 b2) d -> b1 b2 d', {'b1': 12}), \n",
     ")\n",
     "```"
    ]