Weight Normalization for multi-gpu

[MokkeMeguru]

Describe the feature and the current behavior/state.

m_init, v_init = tf.nn.moments(x_init, data_norm_axes)

to

ctx = tf.distribute.get_replica_context()
n = ctx.num_replicas_in_sync

m_init = tf.reduce_mean(x_init, axis=data_norm_axes)
v_init = tf.reduce_mean(tf.square(x_init), axis=data_norm_axes)
m_init, v_init = ctx.all_reduce(
    tf.distribute.ReduceOp.SUM,
    [m_init, v_init]
)

m_init = m_init / n
v_init = v_init / n

Relevant information

• Are you willing to contribute it (yes/no): yes
• Are you willing to maintain it going forward? (yes/no): yes
• Is there a relevant academic paper? (if so, where): no
• Is there already an implementation in another framework? (if so, where): no
• Was it part of tf.contrib? (if so, where): no

Which API type would this fall under (layer, metric, optimizer, etc.)
Weght Normalization layer
https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/layers/wrappers.py#L191

Who will benefit with this feature?
Someone wants to train a large model.

Any other info.
Your data-dependent initialization is so GREAT, Thanks

[WindQAQ]

@MokkeMeguru This is a fair change :-) Feel free to open a PR for it. Thanks.

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/layers/normalization_v2.py#L32

[seanpmorgan]

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