# Path: weakly_supervised_control/disentanglement/tensorsketch/normalization/base.py
# class Norm(Module):
#     """Abstract class for normalization modules."""
# 
#     NAME = "norm"
# 
#     @staticmethod
#     def normalize(*inputs):
#         raise NotImplementedError("Implement normalize.")
# 
#     @staticmethod
#     def add(module, norm, use_out_hook=True):
#         hooks = module.select_hooks_dict(use_out_hook)
#         assert not hasattr(module, norm.NAME), (
#             "{} already in module attributes".format(norm.NAME))
#         assert norm.NAME not in hooks, (
#             "{} already in module hooks".format(norm.NAME))
# 
#         def hook(self, *inputs):
#             return norm(*inputs)
# 
#         setattr(module, norm.NAME, norm)
#         hooks.update({norm.NAME: hook})
# 
#     @staticmethod
#     def remove(module, norm_name, use_out_hook):
#         norm = getattr(module, norm_name)
#         hooks = module.select_hooks_dict(use_out_hook)
# 
#         delattr(module, norm.NAME)
#         del hooks[norm.NAME]
# 
# Path: weakly_supervised_control/disentanglement/tensorsketch/modules/base.py
# def build_with_name_scope(build_parameters):
#     @tf.Module.with_name_scope
#     def build_params_once_with_ns(self, *args):
#         assert not self.built, "{}.built already True".format(self.name)
#         build_parameters(self, *args)
#         self.built = True
#     return build_params_once_with_ns
# 
# Path: weakly_supervised_control/disentanglement/tensorsketch/utils.py
# def assign_moving_average(target, value, momentum):
#     target.assign(momentum * target + (1 - momentum) * value)


import tensorflow as tf
from weakly_supervised_control.disentanglement.tensorsketch.normalization.base import Norm
from weakly_supervised_control.disentanglement.tensorsketch.modules.base import build_with_name_scope
from weakly_supervised_control.disentanglement.tensorsketch.utils import assign_moving_average
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#  This file was modified from `https://github.com/google-research/google-research/blob/master/weak_disentangle`.
"""Batch Normalization utilities."""




class BatchNorm(Norm):
    """Batch Normalization class.
    """

    NAME = "batch_norm"

    def __init__(self, affine=True, momentum=0.9, epsilon=1e-5,
                 track_statistics=True, fuse=None, use_out_hook=None, name=None):
        super().__init__(name=name)
        self.affine = affine
        self.track_statistics = track_statistics
        self.scale = None
        self.bias = None
        self.running_mean = None
        self.running_variance = None
        self.momentum = momentum
        self.epsilon = epsilon
        self.fuse = fuse
        self.use_out_hook = use_out_hook


# Path: weakly_supervised_control/disentanglement/tensorsketch/normalization/base.py
# class Norm(Module):
#     """Abstract class for normalization modules."""
# 
#     NAME = "norm"
# 
#     @staticmethod
#     def normalize(*inputs):
#         raise NotImplementedError("Implement normalize.")
# 
#     @staticmethod
#     def add(module, norm, use_out_hook=True):
#         hooks = module.select_hooks_dict(use_out_hook)
#         assert not hasattr(module, norm.NAME), (
#             "{} already in module attributes".format(norm.NAME))
#         assert norm.NAME not in hooks, (
#             "{} already in module hooks".format(norm.NAME))
# 
#         def hook(self, *inputs):
#             return norm(*inputs)
# 
#         setattr(module, norm.NAME, norm)
#         hooks.update({norm.NAME: hook})
# 
#     @staticmethod
#     def remove(module, norm_name, use_out_hook):
#         norm = getattr(module, norm_name)
#         hooks = module.select_hooks_dict(use_out_hook)
# 
#         delattr(module, norm.NAME)
#         del hooks[norm.NAME]
# 
# Path: weakly_supervised_control/disentanglement/tensorsketch/modules/base.py
# def build_with_name_scope(build_parameters):
#     @tf.Module.with_name_scope
#     def build_params_once_with_ns(self, *args):
#         assert not self.built, "{}.built already True".format(self.name)
#         build_parameters(self, *args)
#         self.built = True
#     return build_params_once_with_ns
# 
# Path: weakly_supervised_control/disentanglement/tensorsketch/utils.py
# def assign_moving_average(target, value, momentum):
#     target.assign(momentum * target + (1 - momentum) * value)


# Path: weakly_supervised_control/disentanglement/tensorsketch/normalization/batch_normalization.py
import tensorflow as tf
from weakly_supervised_control.disentanglement.tensorsketch.normalization.base import Norm
from weakly_supervised_control.disentanglement.tensorsketch.modules.base import build_with_name_scope
from weakly_supervised_control.disentanglement.tensorsketch.utils import assign_moving_average

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#  This file was modified from `https://github.com/google-research/google-research/blob/master/weak_disentangle`.
"""Batch Normalization utilities."""




class BatchNorm(Norm):
    """Batch Normalization class.
    """

    NAME = "batch_norm"

    def __init__(self, affine=True, momentum=0.9, epsilon=1e-5,
                 track_statistics=True, fuse=None, use_out_hook=None, name=None):
        super().__init__(name=name)
        self.affine = affine
        self.track_statistics = track_statistics
        self.scale = None
        self.bias = None
        self.running_mean = None
        self.running_variance = None
        self.momentum = momentum
        self.epsilon = epsilon
        self.fuse = fuse
        self.use_out_hook = use_out_hook


    @build_with_name_scope