Merge pull request #3540 from google:linen-fix-scan-out-axes

Flax Authors · Flax Authors · commit 870847ed3947 · 2023-12-07T16:31:33.000-08:00
PiperOrigin-RevId: 588945218
diff --git a/flax/core/axes_scan.py b/flax/core/axes_scan.py
@@ -100,7 +100,7 @@ def transpose_from_front(ax, xs):
 
     def trans(x):
       if ax < 0:
-        pax = x.ndim - ax
+        pax = x.ndim + ax
       else:
         pax = ax
       assert pax < x.ndim
