test_linalg.py 文件源码

def run_test_matmul_aa_correlator_kernel(self, ntime, nstand, nchan, misalign=0):
        x_shape = (ntime, nchan, nstand*2)
        perm = [1,0,2]
        x8 = ((np.random.random(size=x_shape+(2,))*2-1)*127).astype(np.int8)
        x = x8.astype(np.float32).view(np.complex64).reshape(x_shape)
        x = x.transpose(perm)
        x = x[..., misalign:]
        b_gold = np.matmul(H(x), x)
        triu = np.triu_indices(x.shape[-1], 1)
        b_gold[..., triu[0], triu[1]] = 0
        x = x8.view(bf.DataType.ci8).reshape(x_shape)
        x = bf.asarray(x, space='cuda')
        x = x.transpose(perm)
        x = x[..., misalign:]
        b = bf.zeros_like(b_gold, space='cuda')
        self.linalg.matmul(1, None, x, 0, b)
        b = b.copy('system')
        np.testing.assert_allclose(b, b_gold, RTOL*10, ATOL)