test_multiarray.py 文件源码

def test_choose(self):
        x = 2*np.ones((3,), dtype=int)
        y = 3*np.ones((3,), dtype=int)
        x2 = 2*np.ones((2, 3), dtype=int)
        y2 = 3*np.ones((2, 3), dtype=int)
        ind = np.array([0, 0, 1])

        A = ind.choose((x, y))
        assert_equal(A, [2, 2, 3])

        A = ind.choose((x2, y2))
        assert_equal(A, [[2, 2, 3], [2, 2, 3]])

        A = ind.choose((x, y2))
        assert_equal(A, [[2, 2, 3], [2, 2, 3]])