FSMP 的 unit test 不通过

!14 (merged) 引入的 unit test 失败

python3 fsmp_unit_test.py
Testing combine...
Traceback (most recent call last):
  File "/mnt/stage/xmass/unit/fsmp_unit_test.py", line 71, in <module>
    assert_allclose(cp.asnumpy(test_A_vec), expect_A_vec, **tol_args)
  File "/opt/gentoo/usr/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 1530, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/opt/gentoo/usr/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 844, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=0.002, atol=0

Mismatched elements: 2540 / 2540 (100%)
Max absolute difference: 0.97068812
Max relative difference: 1276.00066514
 x: array([[[0.502443, 0.256718, 0.709304, ..., 0.154647, 0.331047,
         0.153364],
        [0.154269, 0.53244 , 0.688778, ..., 0.853744, 0.080753,...
 y: array([[[ 0.174087, -0.137861,  0.010263, ..., -0.349548,  0.125147,
         -0.335309],
        [-0.174087,  0.137861, -0.010263, ...,  0.349548, -0.125147,...

743715a5 里 combine 并没有失败

$ python3 fsmp_unit_test.py                                                                             
Testing combine...  
Testing move1...                                                                                          
Testing move2...                           
Traceback (most recent call last):                                                                                                                                                                                  
  File "/mnt/stage/xmass/unit/fsmp_unit_test.py", line 132, in <module>
    assert_allclose(cp.asnumpy(test_delta_cx), expect_delta_cx, **tol_args)
  File "/opt/gentoo/usr/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 1530, in assert_allclose
    assert_array_compare(compare, actual, desired, err_msg=str(err_msg),
  File "/opt/gentoo/usr/lib/python3.9/site-packages/numpy/testing/_private/utils.py", line 844, in assert_array_compare
    raise AssertionError(msg)
AssertionError: 
Not equal to tolerance rtol=0.002, atol=0

Mismatched elements: 14905 / 14905 (100%)
Max absolute difference: 4.98417685e-07
Max relative difference: 2.31093733e+15
 x: array([[[-2.499800e-09, -7.908092e-08, -8.237138e-08, ...,
         -1.424661e-07, -9.838610e-08, -6.126912e-08],
        [-1.660199e-09, -3.246853e-09, -3.291737e-09, ...,...
 y: array([[[-9.977718e-21, -8.479367e-19,  2.442424e-18, ...,
          2.179538e-19, -5.283788e-19,  7.807911e-20],
        [ 8.412529e-20,  4.122560e-18,  1.148713e-18, ...,...