mean only and STD only batch normalization - Tensorflow 2

bn_variant.py

from tensorflow.keras import layers
import tensorflow as tf

class MeanOnlyBN(layers.BatchNormalization):
    def __init__(self, **kwargs):
        kwargs['scale'] = False
        kwargs['center'] = False
        kwargs['fused'] = False
        super().__init__(**kwargs)

    def _calculate_mean_and_var(self, inputs, reduction_axes, keep_dims):
        return tf.reduce_mean(inputs, axis=reduction_axes, keepdims=keep_dims), 1


class StdOnlyBN(layers.BatchNormalization):
    def __init__(self, **kwargs):
        kwargs['scale'] =  False
        kwargs['center'] = False
        kwargs['fused'] = False
        super().__init__(**kwargs)

    def _calculate_mean_and_var(self, inputs, reduction_axes, keep_dims):
        return 0, tf.math.reduce_variance(inputs, axis=reduction_axes, keepdims=keep_dims)