def test_matmat_stride0():
a = np.random.random((10, 3, 5))
b = np.random.random((5, 3))
r1 = matmat(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, 3))
r1 = matmat(a, b)
r2 = np.array([np.dot(a, bi) for bi in b])
np.testing.assert_allclose(r1, r2)
def test_matmat_single():
a = np.random.random((3, 5))
b = np.random.random((5, 3))
r1 = matmat(a, b)
r2 = np.dot(a, b)
assert r1.shape == r2.shape
assert r1.dtype == r2.dtype
np.testing.assert_allclose(r1, r2)
def test_matmat_broadcast():
a = np.random.random((10, 11, 2, 4))
b = np.random.random((11, 4, 3))
r1 = matmat(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((5, 3, 2))
r1 = matmat(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_matmat_noncontiguous():
# a non-contiguous
a = np.random.random((10, 3, 3)).swapaxes(1, 2)
b = np.random.random((10, 3, 3))
assert not a.flags.contiguous
r1 = matmat(a, b)
r2 = np.array(list(map(np.dot, a, b)))
np.testing.assert_allclose(r1, r2)
# b non-contiguous
a = np.random.random((10, 3, 4))
b = np.random.random((10, 3, 4)).swapaxes(1, 2)
assert not b.flags.contiguous
r1 = matmat(a, b)
r2 = np.array(list(map(np.dot, a, b)))
np.testing.assert_allclose(r1, r2)
# both non-contiguous
a = np.random.random((10, 3, 5)).swapaxes(1, 2)
b = np.random.random((10, 6, 3)).swapaxes(1, 2)
assert not a.flags.contiguous
assert not b.flags.contiguous
r1 = matmat(a, b)
r2 = np.array(list(map(np.dot, a, b)))
np.testing.assert_allclose(r1, r2)
def test_matmat_ndim():
a = np.random.random((10, 11, 2, 4))
b = np.random.random((10, 11, 4, 3))
r1 = matmat(a, b)
r2 = np.array(list(map(np.dot, a.reshape(-1, 2, 4), b.reshape(-1, 4, 3)))).reshape(10, 11, 2, 3)
np.testing.assert_allclose(r1, r2)
def test_matmat_eqshape():
a = np.random.random((31, 4, 4))
b = np.random.random((31, 4, 4))
r1 = matmat(a, b)
r2 = np.array(list(map(np.dot, a, b)))
np.testing.assert_allclose(r1, r2)
def t_matmat():
matmat(a, b)