Perform shape checks for self.param AFTER unboxing

[danielwatson6]

This pull request will unbox params before performing shape checks in flax/core/scope.py, rather than after.

This should not affect pjit / jax.jit users as unboxing generally doesn't touch traced array data (and wasn't designed for it). However, users that prefer per-device code and manual collectives will run into trouble when implementing certain patterns like FSDP through their own boxes (i.e., subclassing nn.meta.AxisMetadata to achieve a "delayed all-gather") because the shape check fixed here happens at the wrong time.

Fixes # (issue) N/A

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  discussion (please add a
  link).
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.
  (No quality testing = no merge!)

All existing tests pass.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 4 lines in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (31adb00) to head (bc68908).
Report is 87 commits behind head on main.

Files	Patch %	Lines
flax/core/scope.py	0.00%	4 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff           @@
##            main   #4079    +/-   ##
======================================
  Coverage   0.00%   0.00%            
======================================
  Files        106     107     +1     
  Lines      13582   13820   +238     
======================================
- Misses     13582   13820   +238     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.