From 9aa8ec6810da27bf656d22020cc10f9a1bb7e1ae Mon Sep 17 00:00:00 2001
From: Cristian Garcia <cgarcia.e88@gmail.com>
Date: Thu, 7 Dec 2023 23:54:30 +0000
Subject: [PATCH] fix scan out_axes

---
 flax/core/axes_scan.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flax/core/axes_scan.py b/flax/core/axes_scan.py
index 8c51ca3d..2ffd347d 100644
--- 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