From 8112f83064087474c8d32a154f56388a61c97584 Mon Sep 17 00:00:00 2001 From: eadadi <36109667+eadadi@users.noreply.github.com> Date: Sat, 21 Sep 2024 22:22:59 +0300 Subject: [PATCH] Add dict syntax to the einops layers example (#342) * Add dict syntax to the einops layers example * Update 2-einops-for-deep-learning.ipynb * Update 2-einops-for-deep-learning.ipynb * Update docs/2-einops-for-deep-learning.ipynb Co-authored-by: Alex Rogozhnikov --------- Co-authored-by: Alex Rogozhnikov --- docs/2-einops-for-deep-learning.ipynb | 2 ++ 1 file changed, 2 insertions(+) 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", "```" ]