def layer_norm_grad_test()

in blocksparse/norms.py [0:0]

• 36 lines of code
• 5 McCabe index (conditional complexity)

def layer_norm_grad_test(dy, x, g, b, axis=1, segments=1, epsilon=1e-6, relu=False):

    x_shape = x.shape
    K = x_shape[axis]
    if axis == 0:
        dy = dy.reshape(K,-1)
        x  =  x.reshape(K,-1)
        g  =  g.reshape(K, 1)
        b  =  b.reshape(K, 1)
    else:
        axis = 1
        dy = dy.reshape(-1, K)
        x  =  x.reshape(-1, K)
        g  =  g.reshape( 1, K)
        b  =  b.reshape( 1, K)
    K //= segments

    dy = dy.copy()
    dx = np.empty_like(dy)
    dg = np.empty_like(g)
    db = np.empty_like(b)
    for s in range(segments):

        segK  = slice(s*K, s*K+K)
        seg   = (segK, slice(None)) if axis == 0 else (slice(None), segK)

        mean  = np.mean(x[seg], axis=axis, keepdims=True)
        xmean = x[seg] - mean
        xvar  = np.var(x[seg], axis=axis, keepdims=True)
        xstdr = np.reciprocal(np.sqrt(xvar + epsilon))
        xhat  = xmean * xstdr

        if relu:
            dy[seg] = dy[seg] * ((xhat*g[seg] + b[seg]) > 0.0)

        #print("x:%.2f, mean:%.2f, rstd:%.2f, xhat:%.2f, dy:%.2f\n" % (x[0,0], mean[0,0], xstdr[0,0], xhat[0,0], dy[0,0]));

        dg[seg] = np.sum(dy[seg] * xhat, axis=1-axis, keepdims=True)
        db[seg] = np.sum(dy[seg],        axis=1-axis, keepdims=True)
        dy[seg] = dy[seg] * g[seg]

        sum1 = np.sum(xhat * dy[seg], axis=axis, keepdims=True)
        sum2 = np.sum(dy[seg],        axis=axis, keepdims=True)
        dx[seg]   = (dy[seg] - ((xhat * sum1 + sum2) / float(K))) * xstdr

    return dx.reshape(x_shape), dg, db