Merge pull request #2182 from yechengxi:patch-1

PiperOrigin-RevId: 454138229

Author: Flax Authors

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,