def layer_norm()

in blocksparse/norms.py [0:0]

• 24 lines of code
• 11 McCabe index (conditional complexity)

def layer_norm(x, g, b, axis=1, segments=1, epsilon=1e-6, relu=False, atomics=True, bench=0, use_tf=False):

    dev = g.op.device.lower()
    if use_tf or not dev or "cpu" in dev:

        if axis < 0:
            axis += len(x.shape)

        K = x.shape[axis].value
        assert g.shape.num_elements() == K
        assert b.shape.num_elements() == K
        assert K % segments == 0
        assert axis != 0 or segments == 1, "Segments only implemented on axis=1 for now"
        K //= segments

        ys = list()
        for s in range(segments):

            segK = slice(s*K, s*K+K)
            segX = [segK if d == axis else slice(None) for d in range(x.shape.ndims)]

            mean, var = tf.nn.moments(x[segX], [axis], keep_dims=True)
            # mean = tf.reduce_mean(x[segX], axis=[axis], keepdims=True)
            # var  = tf.reduce_mean(tf.square(x[segX] - mean), axis=[axis], keepdims=True)
            norm = (x[segX] - mean) * tf.rsqrt(var + epsilon)
            ys.append(norm * g[segK] + b[segK])

        y = tf.concat(ys, axis) if segments > 1 else ys[0]
        if relu:
            y = tf.nn.relu(y)
    else:
        y, m, v, _, _ = layer_norm_op(x, g, b, S=segments, axis=axis, epsilon=epsilon, relu=relu, atomics=atomics, bench=bench)

    return y