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])


