in src/model.py [0:0]
def norm(x, scope, *, axis=-1, epsilon=1e-5):
"""Normalize to mean = 0, std = 1, then do a diagonal affine transform."""
with tf.variable_scope(scope):
n_state = x.shape[axis].value
g = tf.get_variable('g', [n_state], initializer=tf.constant_initializer(1))
s = tf.reduce_mean(tf.square(x), axis=axis, keepdims=True)
x = x * tf.rsqrt(s + epsilon)
x = x*g
return x