Repository: google/flax

Issue:
ScopeParamShapeError incorrect order of value and initializer shapes
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: jax==0.4.5, jaxlib==0.4.4, flax==0.6.6
- Python version: 3.10
- GPU/TPU model and memory: N/A
- CUDA version (if applicable): N/A

### Problem you have encountered:
I was getting a `ScopeParamShapeError` and it took me awhile to figure out my bug because I got confused by the order of the shapes in the error message. Based on both the wording of the error and the source code, it is implied that the value shape comes first and the initializer shape second:

```
'Inconsistent shapes between value and initializer for parameter "{param_name}" in "{scope_path}": {value_shape}, {init_shape}.'
```

However, based on this minimal example, it seems that the opposite happens:
```python
import jax
import jax.numpy as jnp
import flax.linen as nn

class Model(nn.Module):
    @nn.compact
    def __call__(self, x):
        return nn.Dense(1)(x)

model = Model()
params = model.init(jax.random.PRNGKey(0), jnp.ones(5))["params"]
model.apply({"params": params}, jnp.ones(3))

# ScopeParamShapeError: Inconsistent shapes between value and initializer for parameter "kernel" in "/Dense_0": (5, 1), (3, 1).
```

Unless I've got this completely wrong somehow, it really seems like `(5, 1)` is the initializer shape and `(3, 1)` is the value shape.



Discussion:
Generally all Flax layers are "batch agnostic", meaning you can pass 0 or more batch dimensions. If it helps, imagine you are passing the shapes `(1, 5)` and `(1, 3)`, Dense only cares about the last dimensionh
> Generally all Flax layers are "batch agnostic", meaning you can pass 0 or more batch dimensions. If it helps, imagine you are passing the shapes `(1, 5)` and `(1, 3)`, Dense only cares about the last dimensionh

Yes, but I'm passing 0 batch dimensions in both cases. The dense layer kernel has shape `(n, 1)` where `n` is determined at initialization time by looking at the last dimension of the example input. In this example, `n=5` at initialization time and `n=3` at "apply" time, but the error message reports the opposite.
After looking at the code, you are right! It's reporting the opposite. I think we should improve the error message altogether to make it very clear what shape is expected and what shape was produced.
@chiamp you want to take a loot at this?

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -833,7 +833,7 @@ def param(self, name: str, init_fn: Callable[..., T], *init_args,
         # for inference to a half float type for example.
         if jnp.shape(val) != jnp.shape(abs_val):
           raise errors.ScopeParamShapeError(name, self.path_text,
-                                            jnp.shape(val), jnp.shape(abs_val))
+                                            jnp.shape(abs_val), jnp.shape(val))
     else:
       if not self.is_mutable_collection('params'):
         if self.is_collection_empty('params'):
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -260,9 +260,9 @@ def __call__(self, x):
   """
 
   def __init__(self, param_name, scope_path, value_shape, init_shape):
-    super().__init__('Inconsistent shapes between value and initializer '
-                     f'for parameter "{param_name}" in "{scope_path}": '
-                     f'{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}".')
 
 
 class ScopeVariableNotFoundError(FlaxError):


Test patch:
diff --git a/tests/core/core_scope_test.py b/tests/core/core_scope_test.py
--- a/tests/core/core_scope_test.py
+++ b/tests/core/core_scope_test.py
@@ -107,7 +107,7 @@ def test_inconsistent_param_shapes(self):
     def f(scope):
       scope.param('test', nn.initializers.ones_init(), (4,))
 
-    msg = r'Inconsistent shapes between value and initializer for parameter "test" in "/": \(2,\), \(4,\).'
+    msg = r'Initializer expected to generate shape \(2,\) but got shape \(4,\) instead for parameter "test" in "/"'
     with self.assertRaisesRegex(errors.ScopeParamShapeError, msg):
       apply(f)(freeze({'params': {'test': np.ones((2,))}}))