import tensorflow as tf



def nin(x, num_units, **kwargs):
    s = tf.shape(x)
    sh = x.get_shape().as_list()
    x = tf.reshape(x, [tf.reduce_prod(s[:-1]), sh[-1]])
    x = dense(x, num_units, **kwargs)
    return tf.reshape(x, [-1] + sh[1:-1] + [num_units])


def dense(x, num_units, scope="dense", training=True, ema=None, init=False, bias_initializer=tf.constant_initializer(0.)):
    with tf.variable_scope(scope):
        V = tf.get_variable('V', shape=[int(x.get_shape()[1]), num_units], dtype=tf.float32,
                              initializer=tf.random_normal_initializer(0, 0.05), trainable=True)
        g = tf.get_variable('g', shape=[num_units], dtype=tf.float32,
                              initializer=tf.constant_initializer(1.), trainable=True)
        b = tf.get_variable('b', shape=[num_units], dtype=tf.float32,
                              initializer=bias_initializer, trainable=True)

        def maybe_avg(v):
            if ema is not None and not init:
                v = tf.cond(training, lambda: v, lambda: ema.average(v))
