def sqrt_newton_schulz_autograd(self, A, numIters):
    A_shape = A.shape
    batchSize = A_shape[0]
    dim = A_shape[1]
    normA = (A * A)
    normA = paddle.sum(normA, axis=1)
    normA = paddle.sum(normA, axis=1)
    normA = paddle.sqrt(normA)
    normA1 = normA.reshape([batchSize, 1, 1])
    Y = paddle.divide(A, paddle.expand_as(normA1, A))
    I = paddle.eye(dim, dim).reshape([1, dim, dim])
    l0 = []
    for i in range(batchSize):
        l0.append(I)
    I = paddle.concat(l0, axis=0)
    I.stop_gradient = False
    Z = paddle.eye(dim, dim).reshape([1, dim, dim])
    l1 = []
    for i in range(batchSize):
        l1.append(Z)
    Z = paddle.concat(l1, axis=0)
    Z.stop_gradient = False
    for i in range(numIters):
        T = (0.5 * ((3.0 * I) - Z.bmm(Y)))
        Y = Y.bmm(T)
        Z = T.bmm(Z)
    sA = (Y * paddle.sqrt(normA1).reshape([batchSize, 1, 1]))
    sA = paddle.expand_as(sA, A)
    return sA
def sqrt_newton_schulz_autograd(self, A, numIters):
    A_shape = A.shape
    batchSize = A_shape[0]
    dim = A_shape[1]
    normA = (A * A)
    normA = paddle.sum(normA, axis=1)
    normA = paddle.sum(normA, axis=1)
    normA = paddle.sqrt(normA)
    normA1 = normA.reshape([batchSize, 1, 1])
    Y = paddle.divide(A, paddle.expand_as(normA1, A))
    I = paddle.eye(dim, dim).reshape([1, dim, dim])
    l0 = []
    for i in range(batchSize):
        l0.append(I)
    I = paddle.concat(l0, axis=0)
    I.stop_gradient = False
    Z = paddle.eye(dim, dim).reshape([1, dim, dim])
    l1 = []
    for i in range(batchSize):
        l1.append(Z)
    Z = paddle.concat(l1, axis=0)
    Z.stop_gradient = False
    for i in range(numIters):
        T = (0.5 * ((3.0 * I) - Z.bmm(Y)))
        Y = Y.bmm(T)
        Z = T.bmm(Z)
    sA = (Y * paddle.sqrt(normA1).reshape([batchSize, 1, 1]))
    sA = paddle.expand_as(sA, A)
    return sA