Merge pull request google#4592 from jakevdp:fix-shape

PiperOrigin-RevId: 733430128

Author: Flax Authors

flax/core/scope.py

@@ -956,9 +956,9 @@ def param(
         # NOTE: We could check dtype consistency here as well but it's
         # usefuleness is less obvious. We might intentionally change the dtype
         # for inference to a half float type for example.
-        if jnp.shape(val) != jnp.shape(abs_val):
+        if np.shape(val) != np.shape(abs_val):
           raise errors.ScopeParamShapeError(
-            name, self.path_text, jnp.shape(abs_val), jnp.shape(val)
+            name, self.path_text, np.shape(abs_val), np.shape(val)
           )
     else:
       if not self.is_mutable_collection('params'):

flax/nnx/bridge/module.py

@@ -32,7 +32,7 @@
 from flax.nnx.object import Object
 from flax.nnx import variablelib
 from flax.nnx.bridge import variables as bridge_variables
-import jax.numpy as jnp
+import numpy as np
 
 A = tp.TypeVar('A')
 M = tp.TypeVar('M', bound='Module')
@@ -231,9 +231,9 @@ def param(  # type: ignore[invalid-annotation]
       abs_value_flat = jax.tree_util.tree_leaves(abs_value)
       value_flat = jax.tree_util.tree_leaves(value)
       for val, abs_val in zip(value_flat, abs_value_flat):
-        if jnp.shape(val) != jnp.shape(abs_val):
+        if np.shape(val) != np.shape(abs_val):
           raise errors.ScopeParamShapeError(
-            name, '', jnp.shape(abs_val), jnp.shape(val)
+            name, '', np.shape(abs_val), np.shape(val)
           )
 
       if isinstance(abs_value, variablelib.VariableMetadata):
@@ -282,9 +282,9 @@ def variable(  # type: ignore[invalid-annotation]
       abs_value_flat = jax.tree_util.tree_leaves(abs_value)
       value_flat = jax.tree_util.tree_leaves(value)
       for val, abs_val in zip(value_flat, abs_value_flat):
-        if jnp.shape(val) != jnp.shape(abs_val):
+        if np.shape(val) != np.shape(abs_val):
           raise errors.ScopeParamShapeError(
-            name, '', jnp.shape(abs_val), jnp.shape(val)
+            name, '', np.shape(abs_val), np.shape(val)
           )
 
       if isinstance(abs_value, variablelib.VariableMetadata):