Fix GroupNormalization

[DawyD]

• Fix a bug, where the GroupNormalization layer was normalizing over the second axis instead of the selected axis.
• Update tests (which seem to be irrelevant anyway)

[googlebot]

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

📝 Please visit https://cla.developers.google.com/ to sign.

Once you've signed (or fixed any issues), please reply here with @googlebot I signed it! and we'll verify it.

---

What to do if you already signed the CLA

Individual signers

• It's possible we don't have your GitHub username or you're using a different email address on your commit. Check your existing CLA data and verify that your email is set on your git commits.

Corporate signers

• Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to go/cla#troubleshoot (Public version).
• The email used to register you as an authorized contributor must be the email used for the Git commit. Check your existing CLA data and verify that your email is set on your git commits.
• The email used to register you as an authorized contributor must also be attached to your GitHub account.

ℹ️ Googlers: Go here for more info.

[DawyD]

@googlebot I signed it!

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[seanpmorgan]

Thanks @DawyD!

@Smokrow could you take a look as code owner? We might need add additional test cases

[Smokrow]

As far as I see it the reduction is working as intended.
The axis that is introduced is the group axis and should always be on position 1.

[DawyD]

The problem lies in the reshaping. tf.reshape reads and writes the elements in C-like index order, with the last axis index changing fastest, back to the first axis index changing slowest.

You insert a new axis at index 1 that has size nr_groups and then you reshape assuming that the selected axis will be split into the inserted axis. In your code, this works only if the selected axis is the axis with index 2. If you select, for example, the axis with index 4, the axis that is split into the inserted one is the axis with index 2 and not 4.

[Smokrow]

Hey there. I completely missed this one. Thank you.
@seanpmorgan the fix looks good.

[Smokrow]

@seanpmorgan This will actually change the behavior of the layer since the GroupAxis will be moved to a different position. I am not sure if this can cause trouble down the line. Maybe we should fix that. What do you think?

[googlebot]

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

[seanpmorgan]

Thanks for the fix @DawyD! Apologies for the delayed response, it's been a busy month. As you mentioned the test cases here are only replicating the algorithm in numpy, not verifying against a known expected value. I've created #731 to improve them.

[seanpmorgan]

@seanpmorgan This will actually change the behavior of the layer since the GroupAxis will be moved to a different position. I am not sure if this can cause trouble down the line. Maybe we should fix that. What do you think?

Since the group axis only exists while reshaping in the layer call this shouldn't produce any side-effects as far as I can tell.