Added support for arbitrary tensors for FRN.

[Opletts]

I'll write the test cases, just wanted to confirm if my changes are correct.
@Squadrick @saurabhme
In case of FC Layers with dimensions of B x N, we normalize across N. What's the shape of self.gamma and self.beta supposed to be? I've used 1 x N for now, let me know if I've misunderstood something.

[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.

[bot-of-gabrieldemarmiesse]

@Smokrow

You are owner of some files modified in this pull request.
Would you kindly review the changes whenever you have the time to?
Thank you very much.

[Opletts]

@googlebot I signed it!

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[Squadrick]

@Opletts Thanks for the PR, a few suggested changes.

[Squadrick]

@Opletts The changes look good, can you add the test cases?

[Opletts]

@Squadrick Modified the test cases, limited the tensor sizes to 7 because my system ran out of memory, shouldn't be an issue though.

[Opletts]

@Squadrick @AakashKumarNain just went through the paper again and it doesn't make sense for 2D tensors. FRN is a global normalisation over the spatial extent and that it's done on a per channel basis.

• For 4D tensors with shape = [B x H x W x C], it makes sense that we normalise over the spatial dimensions [H, W].
• For 3D tensors, it HAS to be of form [B x H x C] or [B x C x H], otherwise FRN doesn't make sense.
• 2D tensors have to be of form [H x C] and if it's a single sample, the user should reshape it to [1 x H x C] and then use FRN.

Is my understanding completely flawed? Was I doing the right thing earlier?

[Squadrick]

@Opletts You're right. We assume that the 0th index is batch, so we need 3D at least.

[Opletts]

@Squadrick @AakashKumarNain let me know if everything looks alright, I'll make changes if required.

[seanpmorgan]

Thank you for your contribution. We sincerely apologize for any delay in reviewing, but TensorFlow Addons is transitioning to a minimal maintenance and release mode. New features will not be added to this repository. For more information, please see our public messaging on this decision:
TensorFlow Addons Wind Down

Please consider sending feature requests / contributions to other repositories in the TF community with a similar charters to TFA:
Keras
Keras-CV
Keras-NLP