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

PiperOrigin-RevId: 588945218
google · Dec 8, 2023 · 870847e · 870847e
2 parents 50cd169 + 9aa8ec6
commit 870847e
Showing 1 changed file with 1 addition and 1 deletion.
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