Fix bug when raising ScopeParamNotFoundError.

The order of 'actual value shape' vs. 'expected value shape based on initializer' was mixed up... the actual value should come first and the initialized-based expectation second. This CL changes order to be correct.

This removes a point of confusion when interpreting the error message when raised (expected vs. actual).

PiperOrigin-RevId: 799698584

Author: saugenst

flax/core/scope.py

@@ -948,17 +948,17 @@ def param(
       # catch it with an error message.
       # NOTE: We could consider moving this to `self.`
       abs_value = jax.eval_shape(
-        lambda: init_fn(random.key(0), *init_args, **init_kwargs)
+          lambda: init_fn(random.key(0), *init_args, **init_kwargs)
       )
       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):
-        # 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.
+        # NOTE: We could check dtype consistency here as well but its usefulness
+        # is less obvious. We might intentionally change the dtype for inference
+        # to a half float type for example.
         if np.shape(val) != np.shape(abs_val):
           raise errors.ScopeParamShapeError(
-            name, self.path_text, np.shape(abs_val), np.shape(val)
+              name, self.path_text, np.shape(val), np.shape(abs_val)
           )
     else:
       if not self.is_mutable_collection('params'):

flax/errors.py

@@ -283,9 +283,9 @@ def __call__(self, x):
 
   def __init__(self, param_name, scope_path, value_shape, init_shape):
     super().__init__(
-      f'Initializer expected to generate shape {init_shape} '
-      f'but got shape {value_shape} instead for parameter '
-      f'"{param_name}" in "{scope_path}".'
+        f'For parameter "{param_name}" in "{scope_path}", the given '
+        f'initializer is expected to generate shape {init_shape}, but the '
+        f'existing parameter it received has shape {value_shape}.'
     )
 
 

tests/core/core_scope_test.py

@@ -123,8 +123,9 @@ def f(scope):
       scope.param('test', nn.initializers.ones_init(), (4,))
 
     msg = (
-      r'Initializer expected to generate shape \(2,\) but got shape \(4,\)'
-      r' instead for parameter "test" in "/"'
+        r'For parameter "test" in "/", the given initializer is expected to'
+        r' generate shape \(4,\), but the existing parameter it received has'
+        r' shape \(2,\).'
     )
     with self.assertRaisesRegex(errors.ScopeParamShapeError, msg):
       apply(f)(freeze({'params': {'test': np.ones((2,))}}))