# Path: DeepFried2/utils.py
# def flatten(what, types=(list, tuple), none_to_empty=False):
#     if what is None and none_to_empty:
#         return []
# 
#     if not isinstance(what, types):
#         return [what]
# 
#     # NOTE: I actually timed that this is faster than the comprehension,
#     #       even though it probably doesn't matter :)
#     # 350us vs 250us
#     ret = []
#     for sub in what:
#         ret += flatten(sub, types=types, none_to_empty=none_to_empty)
#     return ret


import DeepFried2 as df
import numpy as _np
from DeepFried2.utils import flatten
from warnings import warn as _warn



class BatchNormalization(df.Module):
    def __init__(self, n_features, momentum=False, eps=1e-5):
        """
        - `n_features` may be an integer (#features, #feature-maps for images) or a tuple.
            - If a single integer, it indicates the size of the 1-axis, i.e. first feature-axis.
              This is the only axis that will be normalized using statistics across all other axes.
            - If a tuple, it indicates the sizes of multiple axes (starting at 1) which are
              considered feature-axes and will consequently be normalized over statistics across all other axes.
        - `momentum` means statistics are collected as running (geometric) statistics
          during training. It should be the decay value within (0,1).
        - `eps` is a small number which is added to the variance in order to
          avoid computing sqrt(0) for features with zero variance.
        """
        df.Module.__init__(self)


# Path: DeepFried2/utils.py
# def flatten(what, types=(list, tuple), none_to_empty=False):
#     if what is None and none_to_empty:
#         return []
# 
#     if not isinstance(what, types):
#         return [what]
# 
#     # NOTE: I actually timed that this is faster than the comprehension,
#     #       even though it probably doesn't matter :)
#     # 350us vs 250us
#     ret = []
#     for sub in what:
#         ret += flatten(sub, types=types, none_to_empty=none_to_empty)
#     return ret


# Path: DeepFried2/layers/BatchNormalization.py
import DeepFried2 as df
import numpy as _np
from DeepFried2.utils import flatten
from warnings import warn as _warn




class BatchNormalization(df.Module):
    def __init__(self, n_features, momentum=False, eps=1e-5):
        """
        - `n_features` may be an integer (#features, #feature-maps for images) or a tuple.
            - If a single integer, it indicates the size of the 1-axis, i.e. first feature-axis.
              This is the only axis that will be normalized using statistics across all other axes.
            - If a tuple, it indicates the sizes of multiple axes (starting at 1) which are
              considered feature-axes and will consequently be normalized over statistics across all other axes.
        - `momentum` means statistics are collected as running (geometric) statistics
          during training. It should be the decay value within (0,1).
        - `eps` is a small number which is added to the variance in order to
          avoid computing sqrt(0) for features with zero variance.
        """
        df.Module.__init__(self)


        self.ndim = len(flatten(n_features))