Update error message for GroupNorm default argument `group_size`

Problem you have encountered:

As default argument group_size is not set to None link the logic here fails. Here’s a small code to repro:

import jax
from jax import numpy as jnp
import flax

rng = jax.random.PRNGKey(0)
gn = flax.linen.GroupNorm(group_size=1)
params = gn.init(rng, jnp.ones((2,2,2)))

What you expected to happen:

Setting group_size alone should be allowed and logic also reflects that. Looks like a small typo to correct though.

Logs, error messages, etc:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-15-3fa6bd6dcb5c> in <module>
      5 rng = jax.random.PRNGKey(0)
      6 gn = flax.linen.GroupNorm(group_size=2)
----> 7 params = gn.init(rng, jnp.ones((2,2,2)))

    [... skipping hidden 11 frame]

/usr/local/lib/python3.8/dist-packages/flax/linen/normalization.py in __call__(self, x)
    320     if ((self.num_groups is None and self.group_size is None) or
    321         (self.num_groups is not None and self.group_size is not None)):
--> 322       raise ValueError('Either `num_groups` or `group_size` should be '
    323                        'specified, but not both of them.')
    324     num_groups = self.num_groups

ValueError: Either `num_groups` or `group_size` should be specified, but not both of them.

Steps to reproduce:

Look at code snippet above.

[marcvanzee]

@sudhakarsingh27 thanks! I think that error should be fixed at head now.

[marcvanzee]

Thanks! For more help see our How to contribute guide