def test_matvec_stride0(): a = np.random.random((10, 3, 5)) b = np.random.random(5) r1 = matvec(a, b) r2 = np.array([np.dot(ai, b) for ai in a]) np.testing.assert_allclose(r1, r2) a = np.random.random((3, 5)) b = np.random.random((12, 5)) r1 = matvec(a, b) r2 = np.array([np.dot(a, bi) for bi in b]) np.testing.assert_allclose(r1, r2)
def test_matvec_single(): a = np.random.random((3, 3)) b = np.random.random(3) r1 = matvec(a, b) r2 = np.dot(a, b) assert r1.shape == r2.shape np.testing.assert_allclose(r1, r2)
def test_matvec_broadcast(): a = np.random.random((10, 31, 4, 5)) b = np.random.random((31, 5)) r1 = matvec(a, b[np.newaxis]) r2 = np.array([ [np.dot(ai, bi) for ai, bi in zip(aj, b)] for aj in a ]) np.testing.assert_allclose(r1, r2) a = np.random.random((7, 6, 3)) b = np.random.random((6, 3)) r1 = matvec(a[:, np.newaxis], b[np.newaxis, :]) r2 = np.array([ [np.dot(ai, bi) for bi in b] for ai in a ]) assert r1.shape == r2.shape np.testing.assert_allclose(r1, r2)
def test_matvec_ndim(): a = np.random.random((10, 31, 4, 5)) b = np.random.random((10, 31, 5)) r1 = matvec(a, b) r2 = np.array(list(map(np.dot, a.reshape(-1, 4, 5), b.reshape(-1, 5)))).reshape(10, 31, 4) np.testing.assert_allclose(r1, r2)
def test_matvec_eqshape(): a = np.random.random((31, 5, 5)) b = np.random.random((31, 5)) r1 = matvec(a, b) r2 = np.array(list(map(np.dot, a, b))) np.testing.assert_allclose(r1, r2)
def t_matvec(): return matvec(a, b)