Merge pull request #2182 from yechengxi:patch-1

PiperOrigin-RevId: 454138229
google · Jun 10, 2022 · 500d349 · 500d349
2 parents c182c7d + 68c47cc
commit 500d349
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flax/core/nn/normalization.py b/flax/core/nn/normalization.py
@@ -20,8 +20,8 @@
 import jax.numpy as jnp
 
 
-def _absolute_dims(rank, dims):
-  return tuple(rank + dim if dim < 0 else dim for dim in dims)
+def _absolute_dims(ndim, dims):
+  return tuple(ndim + dim if dim < 0 else dim for dim in dims)
 
 
 def batch_norm(scope: Scope,