Add GroupNorm to NNX normalization layers

[treigerm]

What does this PR do?

Addresses #4086 and adds a GroupNorm layer to NNX. I tried to follow the Linen implementation and how the other layers have been ported. I also tested equivalence checks between Linen's GroupNorm implementation and the NNX implementation.

Some notes:

• I used num_features instead of num_channels as an input argument to stay consistent with the LayerNorm implementation.
• I deleted the reduced_feature_shape list in the _normalize function as it is unused. If there is a reason why the list was created without being used I can reverse the deletion.

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 GroupNorm missing from NNX normalization layer #4086.
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.
  (No quality testing = no merge!)

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[cgarciae]

Thanks @treigerm! It looks great. I've left a small suggestion.
Can you also please add an entry for GroupNorm here so it renders on the documentation?

[treigerm]

Thanks @cgarciae ! I have added an entry for GroupNorm and also fixed the mypy error.

[cgarciae]

@treigerm can you squash your commits?

[treigerm]

@cgarciae done!

[treigerm]

Had to apply the fix from #4098 to make mypy tests pass locally now.

[codecov-commenter]

Codecov Report

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

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

Files	Patch %	Lines
flax/nnx/nnx/nn/normalization.py	0.00%	46 Missing ⚠️
flax/nnx/__init__.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@          Coverage Diff           @@
##            main   #4095    +/-   ##
======================================
  Coverage   0.00%   0.00%            
======================================
  Files        106     109     +3     
  Lines      13582   14261   +679     
======================================
- Misses     13582   14261   +679     

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