import numpy as np import numpy.linalg as alg import numpy.random as rnd from numpy import * import itertools def identity_like_generalized(a): a = np.array(a) if a.ndim >= 3: r = np.empty(a.shape, dtype=a.dtype) r[...] = np.identity(a.shape[-2]) return r else: return np.identity(a.shape[0]) @test def test_lstsq(): @test def test_rcond(): a = np.array([[0., 1., 0., 1., 2., 0.], [0., 2., 0., 0., 1., 0.], [1., 0., 1., 0., 0., 4.], [0., 0., 0., 2., 3., 0.]]).T b = np.array([1, 0, 0, 0, 0, 0]) x, residuals, rank, s = alg.lstsq(a, b, rcond=-1) assert rank == 4 x, residuals, rank, s = alg.lstsq(a, b) assert rank == 3 x, residuals, rank, s = alg.lstsq(a, b, rcond=None) assert rank == 3 test_rcond() @test def test_empty_a_b(): check = [(4, 2, 2), (0, 4, 1), (0, 4, 2), (4, 0, 1), (4, 0, 2), (4, 2, 0), (0, 0, 0)] i = 0 for values in check: m = values[0] n = values[1] n_rhs = values[2] a = np.arange(m * n).reshape(m, n) b = np.ones((m, n_rhs)) x, residuals, rank, s = alg.lstsq(a, b, rcond=None) if m == 0: i += 1 if i == 1: assert x[0] == 0 elif i == 2: assert x[0][0] == 0 assert np.array_equal(np.array(x.shape), np.array((n, n_rhs))) if m > n: assert np.array_equal(np.array(residuals.shape), np.array((n_rhs, ))) else: assert np.array_equal(np.array(residuals.shape), np.array( (0, ))) if m > n and n_rhs > 0: # residuals are exactly the squared norms of b's columns if n == 0: r = b - np.dot(a, x) assert np.array_equal(residuals, (r * r).sum(axis=-2)) assert rank == np.minimum(m, n) assert np.array_equal(np.array(s.shape), np.array((np.minimum(m, n), ))) test_empty_a_b() test_lstsq() @test def test_matrix_power(): rshft_0 = np.eye(4) rshft_1 = rshft_0[[3, 0, 1, 2]] rshft_2 = rshft_0[[2, 3, 0, 1]] rshft_3 = rshft_0[[1, 2, 3, 0]] rshft_all = [rshft_0, rshft_1, rshft_2, rshft_3] noninv = np.array([[1, 0], [0, 0]]) stacked = np.block([[[rshft_0]]] * 2) # test_large_power rshft = rshft_1.astype(np.float64) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 0), rshft_0) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 1), rshft_1) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 2), rshft_2) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 3), rshft_3) rshft = rshft_1.astype(np.float16) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 0), rshft_0) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 1), rshft_1) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 2), rshft_2) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 3), rshft_3) rshft = rshft_1.astype(np.complex128) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 0), rshft_0) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 1), rshft_1) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 2), rshft_2) assert np.array_equal(alg.matrix_power(rshft, 2**100 + 2**10 + 2**5 + 3), rshft_3) @test def test_power_is_zero(M): mz = alg.matrix_power(M, 0) assert np.array_equal(mz, identity_like_generalized(M)) assert isinstance(mz.dtype, M.dtype) for mat in rshft_all: test_power_is_zero(mat.astype(np.float64)) test_power_is_zero(stacked.astype(np.float64)) for mat in rshft_all: test_power_is_zero(mat.astype(np.float16)) test_power_is_zero(stacked.astype(np.float16)) for mat in rshft_all: test_power_is_zero(mat.astype(np.complex128)) test_power_is_zero(stacked.astype(np.complex128)) @test def test_power_is_one(mat): mz = alg.matrix_power(mat, 1) assert np.array_equal(mz, mat) assert isinstance(mz.dtype, mat.dtype) for mat in rshft_all: test_power_is_one(mat.astype(np.float64)) test_power_is_one(stacked.astype(np.float64)) for mat in rshft_all: test_power_is_one(mat.astype(np.float16)) test_power_is_one(stacked.astype(np.float16)) for mat in rshft_all: test_power_is_one(mat.astype(np.complex128)) test_power_is_one(stacked.astype(np.complex128)) @test def test_power_is_two(mat): mz = alg.matrix_power(mat, 2) assert np.array_equal(mz, np.matmul(mat, mat)) assert isinstance(mz.dtype, mat.dtype) for mat in rshft_all: test_power_is_two(mat.astype(np.float64)) test_power_is_two(stacked.astype(np.float64)) for mat in rshft_all: test_power_is_two(mat.astype(np.float16)) test_power_is_two(stacked.astype(np.float16)) for mat in rshft_all: test_power_is_two(mat.astype(np.complex128)) test_power_is_two(stacked.astype(np.complex128)) @test def test_power_is_minus_one(mat): invmat = alg.matrix_power(mat, -1) assert np.array_equal(identity_like_generalized(mat), np.matmul(invmat, mat)) for mat in rshft_all: test_power_is_minus_one(mat.astype(np.float64)) for mat in rshft_all: test_power_is_minus_one(mat.astype(np.float16)) for mat in rshft_all: test_power_is_minus_one(mat.astype(np.complex128)) test_matrix_power() @test def test_eigvalsh(): @test def test_types(): x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.complex128) w = alg.eigvalsh(x) assert isinstance(w.dtype, np.float64) x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.complex64) w = alg.eigvalsh(x) assert isinstance(w.dtype, np.float32) x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.float64) w = alg.eigvalsh(x) assert isinstance(w.dtype, x.dtype) x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.float32) w = alg.eigvalsh(x) assert isinstance(w.dtype, x.dtype) test_types() @test def test_invalid(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.float32) try: alg.eigvalsh(x, UPLO="lrong") assert False except ValueError: pass try: alg.eigvalsh(x, UPLO="lower") assert False except ValueError: pass try: alg.eigvalsh(x, UPLO="upper") assert False except ValueError: pass test_invalid() @test def test_UPLO(): Klo = np.array([[0, 0], [1, 0]], dtype=np.double) Kup = np.array([[0, 1], [0, 0]], dtype=np.double) tgt = np.array([-1, 1], dtype=np.double) # rtol = get_rtol(np.double) # Check default is 'L' w = alg.eigvalsh(Klo) assert np.array_equal(w, tgt) # Check 'L' w = alg.eigvalsh(Klo, UPLO='L') assert np.array_equal(w, tgt) # Check 'U' w = alg.eigvalsh(Kup, UPLO='U') assert np.array_equal(w, tgt) test_UPLO() @test def test_0_size(): # Check that all kinds of 0-sized arrays work # class ArraySubclass(np.ndarray): # pass a = np.zeros((0, 1, 1), dtype=np.int_) res = alg.eigvalsh(a) assert isinstance(res.dtype, np.float64) assert np.array_equal(np.array((0, 1)), np.array(res.shape)) # This is just for documentation, it might make sense to change: assert isinstance(res, np.ndarray) a = np.zeros((0, 0), dtype=np.complex64) res = alg.eigvalsh(a) assert isinstance(res.dtype, np.float32) assert np.array_equal(np.array((0, )), np.array(res.shape)) # This is just for documentation, it might make sense to change: assert isinstance(res, np.ndarray) test_0_size() test_eigvalsh() @test def test_eigh(): @test def test_types(): x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.complex128) w, v = alg.eigh(x) assert isinstance(w.dtype, np.float64) assert isinstance(v.dtype, np.complex128) x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.complex64) w, v = alg.eigh(x) assert isinstance(w.dtype, np.float32) assert isinstance(v.dtype, np.complex64) x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.float64) w, v = alg.eigh(x) assert isinstance(w.dtype, np.float64) assert isinstance(v.dtype, np.float64) x = np.array([[1., 0.5], [0.5, 1.]], dtype=np.float32) w, v = alg.eigh(x) assert isinstance(w.dtype, np.float32) assert isinstance(v.dtype, np.float32) test_types() @test def test_invalid(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.float32) try: alg.eigh(x, UPLO="lrong") assert False except ValueError: pass try: alg.eigh(x, UPLO="lower") assert False except ValueError: pass try: alg.eigh(x, UPLO="upper") assert False except ValueError: pass test_invalid() @test def test_UPLO(): Klo = np.array([[0, 0], [1, 0]], dtype=np.double) Kup = np.array([[0, 1], [0, 0]], dtype=np.double) tgt = np.array([-1, 1], dtype=np.double) # rtol = get_rtol(np.double) # Check default is 'L' w, v = alg.eigh(Klo) assert np.array_equal(w, tgt) # Check 'L' w, v = alg.eigh(Klo, UPLO='L') assert np.array_equal(w, tgt) # Check 'U' w, v = alg.eigh(Kup, UPLO='U') assert np.array_equal(w, tgt) test_UPLO() @test def test_0_size(): # Check that all kinds of 0-sized arrays work # class ArraySubclass(np.ndarray): # pass a = np.zeros((0, 1, 1), dtype=np.int_) res, res_v = alg.eigh(a) assert isinstance(res_v.dtype, np.float64) assert isinstance(res.dtype, np.float64) assert np.array_equal(np.array(a.shape), np.array(res_v.shape)) assert np.array_equal(np.array((0, 1)), np.array(res.shape)) # This is just for documentation, it might make sense to change: assert isinstance(a, np.ndarray) a = np.zeros((0, 0), dtype=np.complex64) res, res_v = alg.eigh(a) res = alg.eigvalsh(a) assert isinstance(res_v.dtype, np.complex64) assert isinstance(res.dtype, np.float32) assert np.array_equal(np.array(a.shape), np.array(res_v.shape)) assert np.array_equal(np.array((0, )), np.array(res.shape)) # This is just for documentation, it might make sense to change: assert isinstance(a, np.ndarray) test_0_size() test_eigh() @test def test_norm_general(): @test def test_vector_return_type(): a = np.array([1, 0, 1], dtype=np.short) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.ushort) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.intc) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.uintc) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.int_) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.uint) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.longlong) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.ulonglong) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.double) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.longdouble) an = alg.norm(a, -np.inf) assert an == 0.0 a = np.array([1, 0, 1], dtype=np.short) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.ushort) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.intc) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.uintc) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.int_) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.uint) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.longlong) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.ulonglong) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.double) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.longdouble) an = alg.norm(a, 0) assert an == 2 a = np.array([1, 0, 1], dtype=np.short) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.ushort) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.intc) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.uintc) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.int_) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.uint) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.longlong) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.ulonglong) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.double) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.longdouble) an = alg.norm(a, 1) assert an == 2 a = np.array([1, 0, 1], dtype=np.short) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.ushort) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.intc) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.uintc) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.int_) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.uint) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.longlong) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.ulonglong) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.double) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.longdouble) an = alg.norm(a, 2) assert np.isclose(an, 1.41421) a = np.array([1, 0, 1], dtype=np.short) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.ushort) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.intc) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.uintc) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.int_) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.uint) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.longlong) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.ulonglong) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.double) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) a = np.array([1, 0, 1], dtype=np.longdouble) an = alg.norm(a, np.inf) assert np.isclose(an, 1.0) test_vector_return_type() @test def test_vector(): a = [1, 2, 3, 4] b = [-1, -2, -3, -4] c = [-1, 2, -3, 4] @test def _test(v): assert np.isclose(alg.norm(v), 30**0.5) assert np.isclose(alg.norm(v, np.inf), 4.0) assert np.isclose(alg.norm(v, -np.inf), 1.0) assert np.isclose(alg.norm(v, 1), 10.0) assert np.isclose(alg.norm(v, -1), 12.0 / 25) assert np.isclose(alg.norm(v, 2), 30**0.5) assert np.isclose(alg.norm(v, -2), ((205. / 144)**-0.5)) assert np.isclose(alg.norm(v, 0), 4) for v in ( a, b, c, ): _test(v) for v in (np.array(a), np.array(b), np.array(c)): _test(v) test_vector() @test def test_axis(): # Vector norms. # Compare the use of `axis` with computing the norm of each row # or column separately. A = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int_) for order in (None, -1, 0, 1, 2, 3, np.inf, -np.inf): expected0 = [ alg.norm(A[:, k], ord=order) for k in range(A.shape[1]) ] assert np.isclose(alg.norm(A, ord=order, axis=0), expected0).all() expected1 = [ alg.norm(A[k, :], ord=order) for k in range(A.shape[0]) ] assert np.isclose(alg.norm(A, ord=order, axis=1), expected1).all() # Matrix norms. B = np.arange(1, 25, dtype=np.int_).reshape(2, 3, 4) nd = B.ndim for order in (None, -2, 2, -1, 1, np.inf, -np.inf, 'fro'): for axis in itertools.combinations(range(-nd, nd), 2): row_axis, col_axis = axis if row_axis < 0: row_axis += nd if col_axis < 0: col_axis += nd if row_axis == col_axis: try: alg.norm(B, ord=order, axis=axis) assert False except ValueError: pass else: n = alg.norm(B, ord=order, axis=axis) # The logic using k_index only works for nd = 3. # This has to be changed if nd is increased. k_index = nd - (row_axis + col_axis) if row_axis < col_axis: expected = [ alg.norm(B[:].take(k, axis=k_index), ord=order) for k in range(B.shape[k_index]) ] else: expected = [ alg.norm(B[:].take(k, axis=k_index).T, ord=order) for k in range(B.shape[k_index]) ] assert np.isclose(n, expected).all() test_axis() @test def test_keepdims(): A = np.arange(1, 25, dtype=np.int_).reshape(2, 3, 4) # check the order=None, axis=None case expected = alg.norm(A, ord=None, axis=None) found = alg.norm(A, ord=None, axis=None, keepdims=True) assert np.isclose(found[0], expected).all() assert np.array_equal(np.array(found.shape), np.array((1, 1, 1))) # Vector norms. for order in (None, -1, 0, 1, 2, 3, np.inf, -np.inf): for k in range(A.ndim): expected = alg.norm(A, ord=order, axis=k) found = alg.norm(A, ord=order, axis=k, keepdims=True) assert np.isclose(np.squeeze(found, axis=k), expected).all() expected_shape = list(A.shape) expected_shape[k] = 1 assert np.array_equal(np.array(found.shape), np.array(expected_shape)) # Matrix norms. for order in (None, -2, 2, -1, 1, np.inf, -np.inf, 'fro', 'nuc'): for k in itertools.permutations(range(A.ndim), 2): expected = alg.norm(A, ord=order, axis=k) found = alg.norm(A, ord=order, axis=k, keepdims=True) assert np.isclose(np.squeeze(found, axis=k), expected).all() expected_shape = list(A.shape) expected_shape[k[0]] = 1 expected_shape[k[1]] = 1 assert np.array_equal(np.array(found.shape), np.array(expected_shape)) test_keepdims() test_norm_general() @test def test_norm2D(): @test def test_matrix_return_type(): a = np.array([[1, 0, 1], [0, 1, 1]]) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, -np.inf) assert an == 2.0 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, -1) assert np.isclose(an, 1.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, 1) assert an == 2 a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, 2) assert np.isclose(an, 1.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, np.inf) assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, 'fro') assert np.isclose(an, 2.0) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.short) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ushort) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.intc) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uintc) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.int_) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.uint) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longlong) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.ulonglong) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.double) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) a = np.array([[1, 0, 1], [0, 1, 1]], dtype=np.longdouble) an = alg.norm(a, 'nuc') assert np.isclose(an, 2.73205) test_matrix_return_type() @test def test_matrix_2x2(): A = np.array([[1, 3], [5, 7]], dtype=np.int_) assert np.isclose(alg.norm(A), 84**0.5) assert np.isclose(alg.norm(A, 'fro'), 84**0.5) assert np.isclose(alg.norm(A, 'nuc'), 10.0) assert np.isclose(alg.norm(A, np.inf), 12.0) assert np.isclose(alg.norm(A, -np.inf), 4.0) assert np.isclose(alg.norm(A, 1), 10.0) assert np.isclose(alg.norm(A, -1), 6.0) assert np.isclose(alg.norm(A, 2), 9.1231056256176615) assert np.isclose(alg.norm(A, -2), 0.87689437438234041) try: alg.norm(A, 'nofro') assert False except ValueError: pass try: alg.norm(A, -3) assert False except ValueError: pass try: alg.norm(A, 0) assert False except ValueError: pass test_matrix_2x2() @test def test_matrix_3x3(): A = (1 / 10) * \ np.array([[1, 2, 3], [6, 0, 5], [3, 2, 1]], dtype=np.int_) assert np.isclose(alg.norm(A), (1 / 10) * 89**0.5) assert np.isclose(alg.norm(A, 'fro'), (1 / 10) * 89**0.5) assert np.isclose(alg.norm(A, 'nuc'), 1.3366836911774836) assert np.isclose(alg.norm(A, np.inf), 1.1) assert np.isclose(alg.norm(A, -np.inf), 0.6) assert np.isclose(alg.norm(A, 1), 1.0) assert np.isclose(alg.norm(A, -1), 0.4) assert np.isclose(alg.norm(A, 2), 0.88722940323461277) assert np.isclose(alg.norm(A, -2), 0.19456584790481812) test_matrix_3x3() @test def test_bad_args(): A = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int_) B = np.arange(1, 25, dtype=np.int_).reshape(2, 3, 4) for order in [0, 3]: try: alg.norm(A, order, None) assert False except ValueError: pass try: alg.norm(A, order, (0, 1)) assert False except ValueError: pass try: alg.norm(B, order, (1, 2)) assert False except ValueError: pass test_bad_args() test_norm2D() @test def test_norm_nonsystematic(): @test def test_longdouble_norm(): x = np.arange(10, dtype=np.longdouble) assert np.isclose(alg.norm(x, ord=3), 12.6515) test_longdouble_norm() @test def test_intmin(): x = np.array([-2**31], dtype=np.int32) assert np.isclose(alg.norm(x, ord=3), 2147483647.9999974) test_intmin() @test def test_complex_high_ord(): # gh-4156 d = np.empty((2, ), dtype=np.clongdouble) d[0] = 6 + 7j d[1] = -6 + 7j res = 7.55953 assert np.isclose(alg.norm(d, ord=3), res) d = d.astype(np.complex128) assert np.isclose(alg.norm(d, ord=3), res) d = d.astype(np.complex64) assert np.isclose(alg.norm(d, ord=3), res) test_complex_high_ord() test_norm_nonsystematic() @test def test_matrix_rank(): @test def test_matrix_rank(): # Full rank matrix assert 4 == alg.matrix_rank(np.eye(4)) # rank deficient matrix I = np.eye(4) I[-1, -1] = 0. assert alg.matrix_rank(I) == 3 # All zeros - zero rank assert alg.matrix_rank(np.zeros((4, 4))) == 0 # 1 dimension - rank 1 unless all 0 assert alg.matrix_rank([1, 0, 0, 0]) == 1 assert alg.matrix_rank(np.zeros((4, ))) == 0 # accepts array-like assert alg.matrix_rank([1]) == 1 # greater than 2 dimensions treated as stacked matrices ms = np.array((I, np.eye(4), np.zeros((4, 4)))) assert np.array_equal(alg.matrix_rank(ms), np.array([3, 4, 0])) # works on scalar assert alg.matrix_rank(1) == 1 test_matrix_rank() @test def test_symmetric_rank(): assert 4 == alg.matrix_rank(np.eye(4), hermitian=True) assert 1 == alg.matrix_rank(np.ones((4, 4)), hermitian=True) assert 0 == alg.matrix_rank(np.zeros((4, 4)), hermitian=True) # rank deficient matrix I = np.eye(4) I[-1, -1] = 0. assert 3 == alg.matrix_rank(I, hermitian=True) # manually supplied tolerance I[-1, -1] = 1e-8 assert 4 == alg.matrix_rank(I, hermitian=True, tol=0.99e-8) assert 3 == alg.matrix_rank(I, hermitian=True, tol=1.01e-8) test_symmetric_rank() test_matrix_rank() @test def test_reduced_rank(): # Test matrices with reduced rank X = rnd.normal(size=(40, 10)) for i in range(100): # Make a rank deficient matrix X = rnd.normal(size=(40, 10)) X[:, 0] = X[:, 1] + X[:, 2] # Assert that matrix_rank detected deficiency assert alg.matrix_rank(X) == 9 X[:, 3] = X[:, 4] + X[:, 5] assert alg.matrix_rank(X) == 8 test_reduced_rank() @test def test_QR(): @test def check_qr(): a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64) m, n = a.shape k = np.minimum(m, n) # mode == 'complete' res = alg.qr(a, mode='complete') Q, R = res.Q, res.R assert isinstance(Q.dtype, a.dtype) assert isinstance(R.dtype, a.dtype) assert np.array_equal(np.array(Q.shape), np.array((m, m))) assert np.array_equal(np.array(R.shape), np.array((m, n))) assert np.isclose(np.dot(Q, R), a).all() assert np.isclose(np.dot(Q.T.conj(), Q), np.eye(m)).all() assert np.isclose(np.triu(R), R).all() # mode == 'reduced' q1, r1 = alg.qr(a, mode='reduced') assert isinstance(q1.dtype, a.dtype) assert isinstance(r1.dtype, a.dtype) assert np.array_equal(np.array(q1.shape), np.array((m, k))) assert np.array_equal(np.array(r1.shape), np.array((k, n))) assert np.isclose(np.dot(q1, r1), a).all() assert np.isclose(np.dot(q1.T.conj(), q1), np.eye(k)).all() assert np.isclose(np.triu(r1), r1).all() # mode == 'r' r2 = alg.qr(a, mode='r') assert isinstance(r2.dtype, a.dtype) assert np.isclose(r2, r1).all() check_qr() @test def test_qr_empty(): Values = [(3, 0), (0, 3), (0, 0)] for val in Values: m = val[0] n = val[1] k = np.minimum(m, n) a = np.empty((m, n)) h, tau = alg.qr(a, mode='raw') assert isinstance(h.dtype, np.double) assert isinstance(tau.dtype, np.double) assert np.array_equal(np.array(h.shape), np.array((n, m))) assert np.array_equal(np.array(tau.shape), np.array((k, ))) test_qr_empty() @test def test_mode_raw(): # The factorization is not unique and varies between libraries, # so it is not possible to check against known values. Functional # testing is a possibility, but awaits the exposure of more # of the functions in lapack_lite. Consequently, this test is # very limited in scope. Note that the results are in FORTRAN # order, hence the h arrays are transposed. a = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.double) # Test double h, tau = alg.qr(a, mode='raw') assert isinstance(h.dtype, np.double) assert isinstance(tau.dtype, np.double) assert np.array_equal(np.array(h.shape), np.array((2, 3))) assert np.array_equal(np.array(tau.shape), np.array((2, ))) h, tau = alg.qr(a.T, mode='raw') assert isinstance(h.dtype, np.double) assert isinstance(tau.dtype, np.double) assert np.array_equal(np.array(h.shape), np.array((3, 2))) assert np.array_equal(np.array(tau.shape), np.array((2, ))) test_mode_raw() @test def test_mode_all_but_economic(): a = np.array([[1, 2], [3, 4]]) b = np.array([[1, 2], [3, 4], [5, 6]]) @test def check_qr(a): m, n = a.shape k = np.minimum(m, n) # mode == 'complete' res = alg.qr(a, mode='complete') Q, R = res.Q, res.R assert isinstance(Q.dtype, a.dtype) assert isinstance(R.dtype, a.dtype) assert np.array_equal(np.array(Q.shape), np.array((m, m))) assert np.array_equal(np.array(R.shape), np.array((m, n))) assert np.isclose(np.dot(Q, R), a).all() assert np.isclose(np.dot(Q.T.conj(), Q), np.eye(m)).all() assert np.isclose(np.triu(R), R).all() # mode == 'reduced' q1, r1 = alg.qr(a, mode='reduced') assert isinstance(q1.dtype, a.dtype) assert isinstance(r1.dtype, a.dtype) assert np.array_equal(np.array(q1.shape), np.array((m, k))) assert np.array_equal(np.array(r1.shape), np.array((k, n))) assert np.isclose(np.dot(q1, r1), a).all() assert np.isclose(np.dot(q1.T.conj(), q1), np.eye(k)).all() assert np.isclose(np.triu(r1), r1).all() # mode == 'r' r2 = alg.qr(a, mode='r') assert isinstance(r2.dtype, a.dtype) assert np.isclose(r2, r1).all() m1 = a.astype(np.float64) m2 = b.astype(np.float64) check_qr(m1) check_qr(m2) check_qr(m2.T) test_mode_all_but_economic() @test def check_qr_stacked(a): # This test expects the argument `a` to be an ndarray or # a subclass of an ndarray of inexact type. m, n = a.shape[-2:] k = np.minimum(m, n) # mode == 'complete' q, r = alg.qr(a, mode='complete') assert isinstance(q.dtype, a.dtype) assert isinstance(r.dtype, a.dtype) assert np.array_equal(np.array(q.shape[-2:]), np.array((m, m))) assert np.array_equal(np.array(r.shape[-2:]), np.array((m, n))) assert np.isclose(np.matmul(q, r), a).all() I_mat = np.identity(q.shape[-1]) stack_I_mat = np.broadcast_to(I_mat, q.shape[:-2] + (q.shape[-1], ) * 2) assert np.isclose(np.matmul(np.swapaxes(q, -1, -2).conj(), q), stack_I_mat).all() assert np.isclose(np.triu(r[..., :, :]), r).all() # mode == 'reduced' q1, r1 = alg.qr(a, mode='complete') assert isinstance(q1.dtype, a.dtype) assert isinstance(r1.dtype, a.dtype) assert np.array_equal(np.array(q1.shape[-2:]), np.array((m, k))) assert np.array_equal(np.array(r1.shape[-2:]), np.array((k, n))) assert np.isclose(np.matmul(q1, r1), a).all() I_mat = np.identity(q1.shape[-1]) stack_I_mat = np.broadcast_to(I_mat, q1.shape[:-2] + (q1.shape[-1], ) * 2) assert np.isclose(np.matmul(np.swapaxes(q1, -1, -2).conj(), q1), stack_I_mat).all() assert np.isclose(np.triu(r1[..., :, :]), r1).all() # mode == 'r' r2 = alg.qr(a, mode='r') assert isinstance(r2.dtype, a.dtype) assert np.isclose(r2, r1).all() @test def test_stacked_inputs(): # Define values for outer_size & size outer_sizes = ((2, 2), (2, ), (2, 3, 4)) sizes = ((3, 4), (4, 4), (0, 3), (4, 5)) # Iterate over all combinations of outer_size, size, and dt for outer_size in outer_sizes: for size in sizes: # Generate random array A with specified parameters rng = rnd.default_rng(123) A = rng.normal(size=outer_size + size).astype(np.double) B = rng.normal(size=outer_size + size).astype(np.double) check_qr_stacked(A) test_stacked_inputs() test_QR() @test def test_cholesky(): @test def test_basic_property(): shapes = ((1, 1), (2, 2), (3, 3), (3, 10, 10)) for shape in shapes: rnd.seed(1) a = rnd.randn(*shape) t = list(range(len(shape))) t[-2:] = -1, -2 a = np.matmul(a.transpose(t).conj(), a) a = np.asarray(a, dtype=np.float64) c = alg.cholesky(a) b = np.matmul(c, c.transpose(t).conj()) assert np.isclose(a, b).all() d = np.diagonal(c, axis1=-2, axis2=-1) assert np.all(np.isreal(d)) assert np.all(d >= 0) for shape in shapes: rnd.seed(1) a = rnd.randn(*shape) t = list(range(len(shape))) t[-2:] = -1, -2 a = np.matmul(a.transpose(t).conj(), a) a = np.asarray(a, dtype=np.float32) c = alg.cholesky(a) b = np.matmul(c, c.transpose(t).conj()) assert np.isclose(a, b).all() d = np.diagonal(c, axis1=-2, axis2=-1) assert np.all(np.isreal(d)) assert np.all(d >= 0) for shape in shapes: rnd.seed(1) a = rnd.randn(*shape) + [1j] * rnd.randn(*shape) t = list(range(len(shape))) t[-2:] = -1, -2 a = np.matmul(a.transpose(t).conj(), a) a = np.asarray(a, dtype=np.complex64) c = alg.cholesky(a) b = np.matmul(c, c.transpose(t).conj()) # assert np.isclose(a, b).all() d = np.diagonal(c, axis1=-2, axis2=-1) assert np.all(np.isreal(d)) assert np.all(d >= 0) for shape in shapes: rnd.seed(1) a = rnd.randn(*shape) + [1j] * rnd.randn(*shape) t = list(range(len(shape))) t[-2:] = -1, -2 a = np.matmul(a.transpose(t).conj(), a) a = np.asarray(a, dtype=np.complex128) c = alg.cholesky(a) b = np.matmul(c, c.transpose(t).conj()) assert np.isclose(a, b).all() d = np.diagonal(c, axis1=-2, axis2=-1) assert np.all(np.isreal(d)) assert np.all(d >= 0) test_basic_property() @test def test_0_size(): a = np.zeros((0, 1, 1), dtype=np.float64) res = alg.cholesky(a) assert isinstance(res.dtype, np.float64) assert isinstance(res, np.ndarray) assert np.array_equal(np.array(a.shape), np.array(res.shape)) a = np.zeros((1, 0, 0), dtype=np.complex64) res = alg.cholesky(a) assert isinstance(res.dtype, np.complex64) assert isinstance(res, np.ndarray) assert np.array_equal(np.array(a.shape), np.array(res.shape)) test_0_size() test_cholesky() @test def test_outer(): arr1 = np.arange(3) arr2 = np.arange(3) expected = np.array([[0, 0, 0], [0, 1, 2], [0, 2, 4]]) assert np.array_equal(alg.outer(arr1, arr2), expected) test_outer() @test def test_generalized_raise_multiloop(): # It should raise an error even if the error doesn't occur in the # last iteration of the ufunc inner loop invertible = np.array([[1, 2], [3, 4]]) non_invertible = np.array([[1, 1], [1, 1]]) x = np.zeros((4, 4, 2, 2))[1::2] x[...] = invertible x[0, 0] = non_invertible try: alg.inv(x) assert False except alg.LinAlgError: pass test_generalized_raise_multiloop() @test def test_multi_dot(): @test def test_basic_function_with_three_arguments(): # multi_dot with three arguments uses a fast hand coded algorithm to # determine the optimal order. Therefore test it separately. A = rnd.random((6, 2)) B = rnd.random((2, 6)) C = rnd.random((6, 2)) assert np.isclose(alg.multi_dot([A, B, C]), A.dot(B).dot(C)).all() assert np.isclose(alg.multi_dot([A, B, C]), np.dot(A, np.dot(B, C))).all() test_basic_function_with_three_arguments() @test def test_basic_function_with_two_arguments(): # separate code path with two arguments A = rnd.random((6, 2)) B = rnd.random((2, 6)) assert np.isclose(alg.multi_dot([A, B]), A.dot(B)).all() assert np.isclose(alg.multi_dot([A, B]), np.dot(A, B)).all() test_basic_function_with_two_arguments() @test def test_basic_function_with_dynamic_programming_optimization(): # multi_dot with four or more arguments uses the dynamic programming # optimization and therefore deserve a separate A = rnd.random((6, 2)) B = rnd.random((2, 6)) C = rnd.random((6, 2)) D = rnd.random((2, 1)) assert np.isclose(alg.multi_dot([A, B, C, D]), A.dot(B).dot(C).dot(D)).all() test_basic_function_with_dynamic_programming_optimization() @test def test_vector_as_first_argument(): # The first argument can be 1-D A1d = rnd.random(2) # 1-D B = rnd.random((2, 6)) C = rnd.random((6, 2)) D = rnd.random((2, 2)) # the result should be 1-D assert np.array_equal(np.array(alg.multi_dot((A1d, B, C, D)).shape), np.array((2, ))) test_vector_as_first_argument() @test def test_vector_as_last_argument(): # The last argument can be 1-D A = rnd.random((6, 2)) B = rnd.random((2, 6)) C = rnd.random((6, 2)) D1d = rnd.random(2) # 1-D # the result should be 1-D assert np.array_equal(np.array(alg.multi_dot((A, B, C, D1d)).shape), np.array((6, ))) test_vector_as_last_argument() @test def test_vector_as_first_and_last_argument(): # The first and last arguments can be 1-D A1d = rnd.random(2) # 1-D B = rnd.random((2, 6)) C = rnd.random((6, 2)) D1d = rnd.random(2) # 1-D arr = np.array(alg.multi_dot((A1d, B, C, D1d))) dim = arr.ndim assert dim == 0 test_vector_as_first_and_last_argument() @test def test_three_arguments_and_out(): # multi_dot with three arguments uses a fast hand coded algorithm to # determine the optimal order. Therefore test it separately. A = rnd.random((6, 2)) B = rnd.random((2, 6)) C = rnd.random((6, 2)) out = np.zeros((6, 2)) ret = alg.multi_dot([A, B, C], out=out) assert np.array_equal(out, ret) assert np.isclose(out, A.dot(B).dot(C)).all() assert np.isclose(out, np.dot(A, np.dot(B, C))).all() test_three_arguments_and_out() @test def test_two_arguments_and_out(): # separate code path with two arguments A = rnd.random((6, 2)) B = rnd.random((2, 6)) out = np.zeros((6, 6)) ret = alg.multi_dot([A, B], out=out) assert np.array_equal(out, ret) assert np.isclose(out, A.dot(B)).all() assert np.isclose(out, np.dot(A, B)).all() test_two_arguments_and_out() @test def test_dynamic_programming_optimization_and_out(): # multi_dot with four or more arguments uses the dynamic programming # optimization and therefore deserve a separate test A = rnd.random((6, 2)) B = rnd.random((2, 6)) C = rnd.random((6, 2)) D = rnd.random((2, 1)) out = np.zeros((6, 1)) ret = alg.multi_dot((A, B, C, D), out=out) assert np.array_equal(out, ret) assert np.isclose(out, A.dot(B).dot(C).dot(D)).all() test_dynamic_programming_optimization_and_out() # @test # def test_dynamic_programming_logic(): # # Test for the dynamic programming part # # This test is directly taken from Cormen page 376. # arrays = (rnd.random((30, 35)), rnd.random( # (35, 15)), rnd.random((15, 5)), rnd.random( # (5, 10)), rnd.random((10, 20)), rnd.random((20, 25))) # m_expected = np.array([[0., 15750., 7875., 9375., 11875., 15125.], # [0., 0., 2625., 4375., 7125., 10500.], # [0., 0., 0., 750., 2500., 5375.], # [0., 0., 0., 0., 1000., 3500.], # [0., 0., 0., 0., 0., 5000.], # [0., 0., 0., 0., 0., 0.]]) # s_expected = np.array( # [[0, 1, 1, 3, 3, 3], [0, 0, 2, 3, 3, 3], [0, 0, 0, 3, 3, 3], # [0, 0, 0, 0, 4, 5], [0, 0, 0, 0, 0, 5], [0, 0, 0, 0, 0, 0]], # dtype=int) # s_expected -= 1 # Cormen uses 1-based index, python does not. # s, m = alg._multi_dot_matrix_chain_order(arrays, return_costs=True) # # Only the upper triangular part (without the diagonal) is interesting. # assert np.isclose(np.triu(s[:-1, 1:]), np.triu(s_expected[:-1, # 1:])).all() # assert np.isclose(np.triu(m), np.triu(m_expected)).all() # test_dynamic_programming_logic() @test def test_too_few_input_arrays(): try: alg.multi_dot([rnd.random((3, 3))]) assert False except ValueError: pass test_too_few_input_arrays() test_multi_dot() @test def test_tensorinv(): @test def test_non_square_handling(): try: arr = np.ones((4, 6, 8, 2)) ind = 2 alg.tensorinv(arr, ind=ind) assert False except alg.LinAlgError: pass try: arr = np.ones((3, 3, 2)) ind = 1 alg.tensorinv(arr, ind=ind) assert False except alg.LinAlgError: pass test_non_square_handling() @test def test_tensorinv_shape(): shape = (4, 6, 8, 3) a = np.reshape(np.eye(24), shape) ind = 2 ainv = alg.tensorinv(a=a, ind=ind) expected = np.roll(a.shape, -ind) actual = ainv.shape assert np.array_equal(np.array(actual), np.array(expected)) shape = (24, 8, 3) a = np.reshape(np.eye(24), shape) ind = 1 ainv = alg.tensorinv(a=a, ind=ind) expected = np.roll(a.shape, -ind) actual = ainv.shape assert np.array_equal(np.array(actual), np.array(expected)) test_tensorinv_shape() @test def test_tensorinv_ind_limit(): try: shape = (4, 6, 8, 3) a = np.reshape(np.eye(24), shape) ind = -2 alg.tensorinv(a=a, ind=ind) assert False except ValueError: pass try: shape = (4, 6, 8, 3) a = np.reshape(np.eye(24), shape) ind = 0 alg.tensorinv(a=a, ind=ind) assert False except ValueError: pass test_tensorinv_ind_limit() @test def test_tensorinv_result(): # mimic a docstring example shape = (24, 8, 3) a = np.reshape(np.eye(24), shape) ainv = alg.tensorinv(a, ind=1) b = np.ones(24) assert np.isclose(np.tensordot(ainv, b, 1), alg.tensorsolve(a, b)).all() test_tensorinv_result() test_tensorinv() @test def test_tensorsolve(): @test def test_non_square_handling(): try: a = np.array([4, 6, 8, 2]) b = np.ones(a.shape[:2]) alg.tensorsolve(a, b, axes=None) assert False except alg.LinAlgError: pass try: a = np.array((3, 3, 2)) b = np.ones(a.shape[:2]) alg.tensorsolve(a, b, axes=(0, 2)) assert False except ValueError: pass test_non_square_handling() @test def test_tensorsolve_result(): shape = (2, 3, 6) a = rnd.randn(*shape) b = np.ones(a.shape[:2]) x = alg.tensorsolve(a, b) axis: Static[int] = 1 assert np.isclose(np.tensordot(a, x, axes=axis), b).all() shape = (3, 4, 4, 3) a = rnd.randn(*shape) b = np.ones(a.shape[:2]) x = alg.tensorsolve(a, b) axis: Static[int] = 2 assert np.isclose(np.tensordot(a, x, axes=axis), b).all() test_tensorsolve_result() test_tensorsolve() @test def test_unsupported_commontype(): # linalg gracefully handles unsupported type arr = np.array([[1, -2], [2, 5]], dtype='float16') # with assert_raises_regex(TypeError, "unsupported in linalg"): # alg.cholesky(arr) #In python this gives type error test_unsupported_commontype() @test def TestSolve_test_type(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.single) assert alg.solve(x, x)[0][0].__class__.__name__ == 'float32' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.double) assert alg.solve(x, x)[0][0].__class__.__name__ == 'float' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.csingle) assert alg.solve(x, x)[0][0].__class__.__name__ == 'complex64' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.cdouble) assert alg.solve(x, x)[0][0].__class__.__name__ == 'complex' TestSolve_test_type() @test def TestSolve_test_0_size(): a = np.arange(8).reshape(2, 2, 2) b = np.arange(6).reshape(1, 2, 3) expected = alg.solve(a, b)[:, 0:0, :] result = alg.solve(a[:, 0:0, 0:0], b[:, 0:0, :]) assert array_equal(result, expected) assert (isinstance(result, np.ndarray)) try: alg.solve(a[:, 0:0, 0:1], b) assert False except alg.LinAlgError: pass try: alg.solve(a, b[:, 0:0, :]) assert False except ValueError: pass b = np.arange(6).reshape(1, 3, 2) try: alg.solve(a, b) assert False except ValueError: pass try: alg.solve(a[0:0], b[0:0]) assert False except ValueError: pass b = np.arange(2) #.reshape(1, 2) expected = alg.solve(a, b)[0:0] result = alg.solve(a[:, 0:0, 0:0], b[0:0]) assert array_equal(result, expected) assert (isinstance(result, np.ndarray)) b = np.arange(3) #.reshape(1, 3) try: alg.solve(a, b) assert False except ValueError: pass try: alg.solve(a[0:0], b[0:0]) assert False except ValueError: pass try: alg.solve(a[:, 0:0, 0:0], b) assert False except ValueError: pass TestSolve_test_0_size() @test def TestSolve_test_0_size_k(): a = np.arange(4).reshape(1, 2, 2) b = np.arange(6).reshape(3, 2, 1) expected = alg.solve(a, b)[:, :, 0:0] result = alg.solve(a, b[:, :, 0:0]) assert array_equal(result, expected) assert (isinstance(result, np.ndarray)) expected = alg.solve(a, b)[:, 0:0, 0:0] result = alg.solve(a[:, 0:0, 0:0], b[:, 0:0, 0:0]) assert array_equal(result, expected) assert (isinstance(result, np.ndarray)) TestSolve_test_0_size_k() @test def TestInv_test_type(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.single) assert alg.inv(x)[0][0].__class__.__name__ == 'float32' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.double) assert alg.inv(x)[0][0].__class__.__name__ == 'float' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.csingle) assert alg.inv(x)[0][0].__class__.__name__ == 'complex64' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.cdouble) assert alg.inv(x)[0][0].__class__.__name__ == 'complex' TestInv_test_type() @test def TestInv_test_0_size(): a = np.zeros((0, 1, 1), dtype=np.int_) res = alg.inv(a) assert isinstance(res.dtype, np.float64) assert (a.shape == res.shape) assert (isinstance(res, np.ndarray)) a = np.zeros((0, 0), dtype=np.complex64) res = alg.inv(a) assert isinstance(res.dtype, np.complex64) assert (a.shape == res.shape) assert (isinstance(res, np.ndarray)) TestInv_test_0_size() @test def TestEigvals_test_type(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.single) assert alg.eigvals(x)[0][0].__class__.__name__ == 'float32' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.double) assert alg.eigvals(x)[0][0].__class__.__name__ == 'float' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.csingle) assert alg.eigvals(x)[0].__class__.__name__ == 'complex64' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.cdouble) assert alg.eigvals(x)[0].__class__.__name__ == 'complex' x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.single) assert alg.eigvals(x)[0].__class__.__name__ == 'complex64' x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.csingle) assert alg.eigvals(x)[0].__class__.__name__ == 'complex64' x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.double) assert alg.eigvals(x)[0].__class__.__name__ == 'complex' x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.cdouble) assert alg.eigvals(x)[0].__class__.__name__ == 'complex' TestEigvals_test_type() @test def TestEigvals_test_0_size(): a = np.zeros((0, 1, 1), dtype=np.int_) res = alg.eigvals(a) assert isinstance(res.dtype, complex) assert ((0, 1) == res.shape) assert (isinstance(res, np.ndarray)) a = np.zeros((0, 0), dtype=np.complex64) res = alg.eigvals(a) assert isinstance(res.dtype, np.complex64) assert ((0, ) == res.shape) assert (isinstance(res, np.ndarray)) TestEigvals_test_0_size() @test def TestEig_test_types(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.single) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex64) assert isinstance(v.dtype, np.complex64) x = np.array([[1, 0.5], [0.5, 1]], dtype=np.csingle) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex64) assert isinstance(v.dtype, np.complex64) x = np.array([[1, 0.5], [0.5, 1]], dtype=np.double) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex) assert isinstance(v.dtype, np.complex) x = np.array([[1, 0.5], [0.5, 1]], dtype=np.cdouble) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex) assert isinstance(v.dtype, np.complex) x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.single) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex64) assert isinstance(v.dtype, np.complex64) x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.csingle) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex64) assert isinstance(v.dtype, np.complex64) x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.double) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex) assert isinstance(v.dtype, np.complex) x = np.array([[1.0, 0.5], [-1.0, 1.0]], dtype=np.cdouble) w, v = alg.eig(x) assert isinstance(w.dtype, np.complex) assert isinstance(v.dtype, np.complex) TestEig_test_types() @test def TestEig_test_0_size(): a = np.zeros((0, 1, 1), dtype=np.int_) res, res_v = alg.eig(a) assert (isinstance(res_v.dtype, np.complex)) assert (isinstance(res.dtype, np.complex)) assert (a.shape == res_v.shape) assert ((0, 1) == res.shape) assert (isinstance(a, np.ndarray)) a = np.zeros((0, 0), dtype=np.complex64) res, res_v = alg.eig(a) assert isinstance(res_v.dtype, np.complex64) assert isinstance(res.dtype, np.complex64) assert (a.shape == res_v.shape) assert ((0, ) == res.shape) assert (isinstance(a, np.ndarray)) TestEig_test_0_size() @test def SVDBaseTests_test_types(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.single) res = alg.svd(x) U, S, Vh = res.U, res.S, res.Vh assert isinstance(U.dtype, np.float32) assert isinstance(S.dtype, np.float32) assert isinstance(Vh.dtype, np.float32) s = alg.svd(x, compute_uv=False) assert isinstance(s.dtype, np.float32) x = np.array([[1, 0.5], [0.5, 1]], dtype=np.csingle) res = alg.svd(x) U, S, Vh = res.U, res.S, res.Vh assert isinstance(U.dtype, np.complex64) assert isinstance(S.dtype, np.float32) assert isinstance(Vh.dtype, np.complex64) s = alg.svd(x, compute_uv=False) assert isinstance(s.dtype, np.float32) x = np.array([[1, 0.5], [0.5, 1]], dtype=np.double) res = alg.svd(x) U, S, Vh = res.U, res.S, res.Vh assert isinstance(U.dtype, np.float) assert isinstance(S.dtype, np.float) assert isinstance(Vh.dtype, np.float) s = alg.svd(x, compute_uv=False) assert isinstance(s.dtype, np.float) x = np.array([[1, 0.5], [0.5, 1]], dtype=np.cdouble) res = alg.svd(x) U, S, Vh = res.U, res.S, res.Vh assert isinstance(U.dtype, np.complex) assert isinstance(S.dtype, np.float) assert isinstance(Vh.dtype, np.complex) s = alg.svd(x, compute_uv=False) assert isinstance(s.dtype, np.float) SVDBaseTests_test_types() @test def TestSVD_test_empty_identity(): x = np.empty((4, 0)) u, s, vh = alg.svd(x, compute_uv=True) assert (u.shape == (4, 4)) assert (vh.shape == (0, 0)) assert (u == np.eye(4)).all() x = np.empty((0, 4)) u, s, vh = alg.svd(x, compute_uv=True) assert (u.shape == (0, 0)) assert (vh.shape == (4, 4)) assert (vh == np.eye(4)).all() TestSVD_test_empty_identity() @test def TestCond_test_basic_nonsvd(): A = np.array([[1, 0, 1], [0, -2, 0], [0, 0, 3]]) assert (alg.cond(A, inf) == 4) assert (alg.cond(A, -inf) == 2 / 3) assert (alg.cond(A, 1) == 4) assert (alg.cond(A, -1) == 0.5) assert (alg.cond(A, 'fro') == np.sqrt(265 / 12)) TestCond_test_basic_nonsvd() @test def TestCond_test_singular(): assert (alg.cond([[0, 0], [0, 0]], None) > 1e15) assert (alg.cond([[0, 0], [0, 0]], 1) > 1e15) assert (alg.cond([[0, 0], [0, 0]], 2) > 1e15) assert (alg.cond([[0, 0], [0, 0]], 'fro') > 1e15) assert (alg.cond([[0, 0], [0, 0]], None) > 1e15) assert (alg.cond([[1, 1], [1, 1]], 1) > 1e15) assert (alg.cond([[1, 1], [1, 1]], 2) > 1e15) assert (alg.cond([[1, 1], [1, 1]], 'fro') > 1e15) TestCond_test_singular() @test def TestCond_test_nan(): ps = (None, 1, -1, 2, -2, 'fro') p_pos = (None, 1, 2, 'fro') A = np.ones((2, 2)) A[0, 1] = np.nan for p in ps: for p in (1, -1, 'fro'): c = alg.cond(A, p) assert (isinstance(c, np.float64)) assert (np.isnan(c)) for p in (None, 2, -2): try: alg.cond(A, p) if not __apple__: # TODO: Accelerate doesn't give an error on SVD(nan) assert False except alg.LinAlgError: pass ps = (1, -1, 'fro') p_pos = (None, 1, 2, 'fro') A = np.ones((3, 2, 2)) A[1, 0, 1] = np.nan for p in ps: c = alg.cond(A, p) assert (np.isnan(c[1])) if p in p_pos: assert (c[0] > 1e15) assert (c[2] > 1e15) else: assert (not np.isnan(c[0])) assert (not np.isnan(c[2])) TestCond_test_nan() @test def TestCond_test_stacked_singular(): rnd.seed(1234) A = rnd.rand(2, 2, 2, 2) A[0, 0] = 0 A[1, 1] = 0 for p in (None, 1, 2, 'fro', -1, -2): c = alg.cond(A, p) assert (c[0, 0] == np.inf) assert (c[1, 1] == np.inf) assert (np.isfinite(c[0, 1])) assert (np.isfinite(c[1, 0])) TestCond_test_stacked_singular() @test def TestDet_test_zero(): assert (alg.det([[0.0]]) == 0.0) assert (isinstance(alg.det([[0.0]]), double)) assert (alg.det([[0.0j]]) == 0.0) assert (isinstance(alg.det([[0.0j]]), cdouble)) assert (alg.slogdet([[0.0]]) == (0.0, -inf)) assert (isinstance(alg.slogdet([[0.0]])[0], double)) assert (isinstance(alg.slogdet([[0.0]])[1], double)) assert (alg.slogdet([[0.0j]]) == (0.0j, -inf)) assert (isinstance(alg.slogdet([[0.0j]])[0], cdouble)) assert (isinstance(alg.slogdet([[0.0j]])[1], double)) TestDet_test_zero() @test def TestDet_test_types(): x = np.array([[1, 0.5], [0.5, 1]], dtype=np.single) ph, s = alg.slogdet(x) assert s.__class__.__name__ == 'float32' assert ph.__class__.__name__ == 'float32' assert alg.det(x).__class__.__name__ == 'float32' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.double) ph, s = alg.slogdet(x) assert s.__class__.__name__ == 'float' assert ph.__class__.__name__ == 'float' assert alg.det(x).__class__.__name__ == 'float' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.csingle) ph, s = alg.slogdet(x) assert s.__class__.__name__ == 'float32' assert ph.__class__.__name__ == 'complex64' assert alg.det(x).__class__.__name__ == 'complex64' x = np.array([[1, 0.5], [0.5, 1]], dtype=np.cdouble) ph, s = alg.slogdet(x) assert s.__class__.__name__ == 'float' assert ph.__class__.__name__ == 'complex' assert alg.det(x).__class__.__name__ == 'complex' TestDet_test_types() @test def TestDet_test_0_size(): a = np.zeros((0, 0), dtype=np.complex64) res = alg.det(a) assert (res == 1.) assert (isinstance(res, np.complex64)) res = alg.slogdet(a) assert ((res[0], res[1]) == (np.complex64(1 + 0j), np.float32(0))) assert (isinstance(res[0], np.complex64)) assert isinstance(res[1], np.float32) a = np.zeros((0, 0), dtype=np.float64) res = alg.det(a) assert (res == 1.) assert isinstance(res, np.float64) res = alg.slogdet(a) assert ((res[0], res[1]) == (1., 0.)) assert isinstance(res[0], np.float64) assert isinstance(res[1], np.float64) TestDet_test_0_size() @test def test_trace(): x = np.arange(60).reshape((3, 4, 5)) actual = alg.trace(x) expected = np.array([75, 78, 81, 84, 87]) assert array_equal(actual, expected) test_trace() @test def test_tensordot(): x = np.arange(6).reshape((2, 3)) assert alg.tensordot(x, x) == 55 assert alg.tensordot(x, x, axes=[(0, 1), (0, 1)]) == 55 test_tensordot() @test def test_matmul(): x = np.arange(6).reshape((2, 3)) actual = alg.matmul(x, x.T) expected = np.array([[5, 14], [14, 50]]) assert array_equal(actual, expected) test_matmul() @test def test_dot(a, b, expected): assert np.dot(a, b) == expected test_dot(3, 4, 12) test_dot(np.array([2j, 3j]), np.array([2j, 3j]), -13 + 0j) test_dot(np.array([1, 2, 3]), np.array([3, 2, 1]), 10) test_dot(np.array([1.5, 2.3, 3.2]), np.array([3, 2, 1]), 12.3) @test def test_dot_multidim(a, b, expected): assert (np.dot(a, b) == expected).all() test_dot_multidim(np.array([[1, 0], [0, 1]]), np.array([[4, 1], [2, 2]]), np.array([[4, 1], [2, 2]])) test_dot_multidim(np.array([[4, 5], [8, 9]]), np.array([2, 3]), np.array([23, 43])) @test def test_multi_dot(arrays, expected): assert (alg.multi_dot(arrays) == expected).all() test_multi_dot([ np.array([[1, 2], [4, 5]]), np.array([[5, 6], [7, 9]]), np.array([[6, 1], [3, 4]]) ], np.array([[186, 115], [537, 331]])) test_multi_dot( (np.array([1, 2]), np.array([[5, 6], [7, 9]]), np.array([[6, 1], [3, 4]])), np.array([186, 115])) @test def test_vdot(a, b, expected): assert np.vdot(a, b) == expected test_vdot(np.array([1 + 2j, 3 + 4j]), np.array([5 + 6j, 7 + 8j]), 70 - 8j) test_vdot(np.array([[1, 4], [5, 6]]), np.array([[4, 1], [2, 2]]), 30) @test def test_inner(a, b, expected): if isinstance(a, int) or isinstance(b, int): assert (np.inner(a, b) == expected).all() elif (staticlen(a.shape) == 1 and staticlen(b.shape) == 1): assert np.inner(a, b) == expected else: assert (np.inner(a, b) == expected).all() test_inner(np.array([1, 2, 3]), np.array([0, 1, 0]), 2) test_inner( np.arange(24).reshape((2, 3, 4)), np.arange(4), np.array([[14, 38, 62], [86, 110, 134]])) test_inner(np.eye(2), 7, np.array([[7., 0.], [0., 7.]])) @test def test_outer(a, b, expected): assert (np.outer(a, b) == expected).all() test_outer(np.array([4, 5, 6]), np.array([1, 2, 3]), np.array([[4, 8, 12], [5, 10, 15], [6, 12, 18]])) test_outer(np.array([[4, 6], [3, 8]]), np.array([1, 2, 3]), np.array([[4, 8, 12], [6, 12, 18], [3, 6, 9], [8, 16, 24]])) @test def test_matmul(x1, x2, expected): assert (np.matmul(x1, x2) == expected).all() test_matmul(np.array([[1, 0], [0, 1]]), np.array([[4, 1], [2, 2]]), np.array([[4, 1], [2, 2]])) test_matmul(np.array([[1, 0], [0, 1]]), np.array([1, 2]), np.array([1, 2])) @test def test_matmul_sign(x1, x2, expected): assert (x1 @ x2) == expected test_matmul_sign(np.array([2j, 3j]), np.array([2j, 3j]), -13 + 0j) @test def test_tensordot(a, b, expected, axes=2): assert (np.tensordot(a, b, axes) == expected).all() test_tensordot(np.arange(60.).reshape(3, 4, 5), np.arange(24.).reshape(4, 3, 2), np.array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]]), axes=([1, 0], [0, 1])) #@test #def test_einsum(subscripts, operands, expected): # assert (np.einsum(subscripts, operands) == expected).all() # #test_einsum('ii->i', np.arange(25).reshape(5,5), np.array([0, 6, 12, 18, 24])) @test def test_matrix_power(a, n, expected): assert (alg.matrix_power(a, n) == expected).all() test_matrix_power(np.array([[0, 1], [-1, 0]]), 3, np.array([[0, -1], [1, 0]])) test_matrix_power(np.array([[0, 1], [-1, 0]]), 0, np.array([[1, 0], [0, 1]])) test_matrix_power(np.array([[0., 1.], [1., 0.]]), -3, np.array([[0., 1.], [1., 0.]])) test_matrix_power( np.array([[0., -1., 0., 0.], [1., 0., 0., 0.], [0., 0., 0., 1.], [0., 0., -1., 0.]]), 2, np.array([[-1., 0., 0., 0.], [0., -1., 0., 0.], [0., 0., -1., 0.], [0., 0., 0., -1.]])) @test def test_kron(a, b, expected): assert (np.kron(a, b) == expected).all() test_kron(np.array([1, 10, 100]), np.array([5, 6, 7]), np.array([5, 6, 7, 50, 60, 70, 500, 600, 700])) test_kron( np.eye(2), np.ones((2, 2)), np.array([[1., 1., 0., 0.], [1., 1., 0., 0.], [0., 0., 1., 1.], [0., 0., 1., 1.]])) test_kron(np.array([1, 10, 100]), np.array([5, 6]), np.array([5, 6, 50, 60, 500, 600])) test_kron(np.array([1, 10]), np.array([5, 6, 7]), np.array([5, 6, 7, 50, 60, 70])) test_kron(10, np.array([5, 6, 7]), np.array([50, 60, 70])) test_kron(np.array([1, 10]), 8.8, np.array([8.8, 88.])) test_kron(np.array([1, 10, 100]), np.array([5, 6, 7.1]), np.array([5., 6., 7.1, 50., 60., 71., 500., 600., 710.])) @test def test_cholesky(a, expected): assert (alg.cholesky(a) == expected).all() test_cholesky(np.array([[1, -2], [2, 5]]), np.array([[1., 0.], [2., 1.]])) test_cholesky(np.array([[1]]), np.array([[1.]])) @test def test_qr(a, expected_Q, expected_R, mode: Static[str] = 'reduced'): Q, R = alg.qr(a, mode=mode) assert (round(Q, 8) == expected_Q).all() and (round(R, 8) == expected_R).all() test_qr(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 0]]), np.array([[-0.12309149, 0.90453403, 0.40824829], [-0.49236596, 0.30151134, -0.81649658], [-0.86164044, -0.30151134, 0.40824829]]), np.array([[-8.1240384, -9.6011363, -3.32347026], [0., 0.90453403, 4.52267017], [0., 0., -3.67423461]]), mode='reduced') test_qr(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 0]]), np.array([[-0.12309149, 0.90453403, 0.40824829], [-0.49236596, 0.30151134, -0.81649658], [-0.86164044, -0.30151134, 0.40824829]]), np.array([[-8.1240384, -9.6011363, -3.32347026], [0., 0.90453403, 4.52267017], [0., 0., -3.67423461]]), mode='complete') test_qr(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 0]]), np.array([[-8.12403840, 0.43840236, 0.76720414], [-9.60113630, 0.90453403, 0.90907633], [-3.32347026, 4.52267017, -3.67423461]]), np.array([1.12309149, 1.09503851, 0.]), mode='raw') gen = rnd.Generator(rnd.PCG64(1234)) a = gen.random((2, 3)) + gen.random((2, 3)) * 1j test_qr( a, np.array([[-0.91062381 - 0.22541026j, 0.33533976 + 0.08661254j], [-0.24398834 - 0.24581330j, -0.65410584 - 0.67245153j]]), np.array([[ -1.07256120 + 0.j, -0.60427723 - 0.23352661j, -1.23677059 - 0.78957732j ], [0. + 0.j, -0.35019471 + 0.j, -0.09425020 - 0.07618002j]])) @test def test_svd(a, full_matrices: bool = True, compute_uv: Static[int] = True, hermitian: bool = False, expected=0): if hermitian: x = alg.svd(a, full_matrices=full_matrices, compute_uv=compute_uv, hermitian=hermitian) assert (round(x[0], 8) == expected).all() if not compute_uv: x = alg.svd(a, full_matrices=full_matrices, compute_uv=compute_uv, hermitian=hermitian) assert (round(x, 8) == expected).all() else: U, S, Vh = alg.svd(a, full_matrices=full_matrices, compute_uv=compute_uv, hermitian=hermitian) if not full_matrices: np.allclose(a, np.matmul(U * S[..., None, :], Vh)) elif full_matrices: np.allclose(a, np.dot(U[:, :len(S)] * S, Vh)) test_svd(np.array([[1, 2], [3, 4], [5, 6], [7, 8]]), full_matrices=False) test_svd(np.array([[1, 2], [3, 4], [5, 6]]), full_matrices=True) test_svd(np.array([[1, 2], [3, 4]]), compute_uv=False, expected=np.array([5.46498570, 0.36596619])) test_svd(np.array([[1, 2], [3, 4]]), hermitian=True, expected=np.array([[0.52573111, -0.85065081], [0.85065081, 0.52573111]])) gen = rnd.Generator(rnd.PCG64(1234)) a = gen.random((9, 6)) + gen.random((9, 6)) * 1j test_svd(a, full_matrices=True) @test def test_eig(a, expected_val, expected_vec): eigenvalues, eigenvectors = alg.eig(a) assert (round(eigenvalues, 8) == expected_val).all() and (round( eigenvectors, 8) == expected_vec).all() test_eig(np.array([[1, -2], [-2, -5]]), np.array([1.60555128, -5.60555128]), np.array([[0.95709203, 0.28978415], [-0.28978415, 0.95709203]])) test_eig(np.array([[1 + 1e-9, 0], [0, 1 - 1e-9]]), np.array([1., 1.]), np.array([[1., 0.], [0., 1.]])) @test def test_eigh(a, expected_val, expected_vec, uplo: str = 'L'): eigenvalues, eigenvectors = alg.eigh(a) assert (round(eigenvalues, 8) == expected_val).all() and (round( eigenvectors, 8) == expected_vec).all() test_eigh(np.array([[1, -2], [2, 5]]), np.array([0.17157288, 5.82842712]), np.array([[-0.92387953, 0.38268343], [0.38268343, 0.92387953]])) test_eigh( np.array([[1 + 0j, -2j], [2j, 5 + 0j]]), np.array([0.17157288, 5.82842712]), np.array([[-0.92387953 + 0.j, -0.38268343 + 0.j], [0. + 0.38268343j, 0. - 0.92387953j]])) @test def test_eigvals(a, expected): assert (round(alg.eigvals(a), 8) == expected).all() test_eigvals(np.array([[9, 8], [6, 7]]), np.array([15., 1.])) test_eigvals(np.array([[1, -2], [-2, -5]]), np.array([1.60555128, -5.60555128])) @test def test_eigvalsh(a, expected): assert (round(alg.eigvalsh(a), 8) == expected).all() test_eigvalsh(np.array([[1, -2], [2, 5]]), np.array([0.17157288, 5.82842712])) test_eigvalsh(np.array([[1 + 0j, -2j], [2j, 5 + 0j]]), np.array([0.17157288, 5.82842712])) @test def test_norm(x, expected, ord=None, axis=None): if axis is None: assert alg.norm(x, ord, axis) == expected else: assert (round(alg.norm(x, ord, axis), 8) == expected).all() test_norm(np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4]), 7.745966692414834) test_norm(np.array([[-4, -3, -2], [-1, 0, 1], [2, 3, 4]]), 7.745966692414834) test_norm(np.array([[-4, -3, -2], [-1, 0, 1], [2, 3, 4]]), 7.745966692414834, 'fro') test_norm(np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4]), 4.0, np.inf) test_norm(np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4]), 0.0, -np.inf) test_norm(np.array([[-4, -3, -2], [-1, 0, 1], [2, 3, 4]]), 9.0, np.inf) test_norm(np.array([[-4, -3, -2], [-1, 0, 1], [2, 3, 4]]), 2.0, -np.inf) test_norm(np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4]), 20.0, 1) test_norm(np.array([[-4, -3, -2], [-1, 0, 1], [2, 3, 4]]), 7.0, 1) test_norm(np.array([[-4, -3, -2], [-1, 0, 1], [2, 3, 4]]), 6.0, -1) test_norm(np.array([[1, 2, 3], [-1, 1, 4]]), np.array([1.41421356, 2.23606798, 5.]), axis=0) test_norm(np.array([[1, 2, 3], [-1, 1, 4]]), np.array([6., 6.]), ord=1, axis=1) @test def test_cond(x, expected, p=None): assert alg.cond(x, p) == expected test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 1.4142135623730951) test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 3.1622776601683795, 'fro') test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 2.0, np.inf) test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 1.0, -np.inf) test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 2.0, 1) test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 1.0, -1) test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 1.4142135623730951, 2) test_cond(np.array([[1, 0, -1], [0, 1, 0], [1, 0, 1]]), 0.70710678118654746, -2) @test def test_det(a, expected): if isinstance(expected, float): assert np.isclose(alg.det(a), expected) else: assert np.allclose(alg.det(a), expected) test_det(np.array([[1, 2], [3, 4]]), -2.0) test_det(np.array([[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]), np.array([-2., -3., -8.])) @test def test_matrix_rank(A, expected, tol=None): assert alg.matrix_rank(A, tol) == expected test_matrix_rank(np.ones((4, )), 1) test_matrix_rank(np.zeros((4, )), 0) test_matrix_rank( np.array([[1., 0., 0., 0.], [0., 1., 0., 0.], [0., 0., 1., 0.], [0., 0., 0., 0.]]), 3) test_matrix_rank(np.eye(4), 4) test_matrix_rank(np.array([[1., 2., 3.], [-1., -2., -3.], [4., 5., 6.], [-4., -5., -6.]]), 1, tol=1.5) @test def test_slogdet(a, expected_sign, expected_log): sign, logabsdet = alg.slogdet(a) if (isinstance(expected_sign, float) or isinstance( expected_sign, int)) and isinstance(expected_log, float): assert sign == expected_sign and round(logabsdet, 9) == expected_log else: assert (sign == expected_sign).all() and (round(logabsdet, 8) == expected_log).all() test_slogdet(np.array([[[1, 2], [3, 4]], [[1, 2], [2, 1]], [[1, 3], [3, 1]]]), np.array([-1., -1., -1.]), np.array([0.69314718, 1.09861229, 2.07944154])) test_slogdet(np.eye(500) * 0.1, 1, -1151.292546497) @test def test_trace(a, expected, offset=0, axis1=0, axis2=1): assert np.asarray(np.trace(a, offset, axis1, axis2) == expected).all() assert np.asarray(a.trace(offset, axis1, axis2) == expected).all() test_trace(np.eye(3), 3.0) test_trace( np.arange(24).reshape((2, 2, 2, 3)), np.array([[18, 20, 22], [24, 26, 28]])) test_trace(np.arange(24).reshape((2, 2, 2, 3)), np.array([[6, 7, 8], [9, 10, 11]]), offset=1) test_trace(np.arange(24).reshape((2, 2, 2, 3)), np.array([[12, 13, 14], [15, 16, 17]]), offset=-1) test_trace(np.arange(24).reshape((2, 2, 2, 3)), np.array([[0, 0, 0], [0, 0, 0]]), offset=2) test_trace(np.arange(24).reshape((2, 2, 2, 3)), np.array([[0, 0, 0], [0, 0, 0]]), offset=-2) test_trace(np.arange(24).reshape((2, 2, 2, 3)), np.array([[9, 11, 13], [33, 35, 37]]), axis1=2) test_trace(np.arange(24).reshape((2, 2, 2, 3)), np.array([[13, 19], [25, 31]]), axis2=3) @test def test_solve(a, b): x = alg.solve(a, b) assert np.allclose(np.dot(a, x), b) test_solve(np.array([[1, 2], [3, 5]]), np.array([1, 2])) @test def test_tensorsolve(a, b, expected, axes=None): x = alg.tensorsolve(a, b, axes) assert (round(x, 8) == expected).all() test_tensorsolve(np.array([[60, 40, 50], [10, 20, 30], [70, 80, 90]]), np.array([18, 19, 20]), np.array([-0.04444444, -1.76111111, 1.82222222])) @test def test_lstsq(a, b, expected_ls, expected_res, rcond=None): ls, res, m, c = alg.lstsq(a, b, rcond) assert (round(ls, 8) == expected_ls).all() and (round(res, 8) == expected_res).all() test_lstsq(np.array([[1, 2], [-1, 1], [0, 3]]), np.array([1, 3, 0]), np.array([-1.22222222, 0.44444444]), np.array([5.33333333])) test_lstsq(np.array([[2, 1], [1, 3]]), np.array([[9, 8], [4, 6]]), round(np.array([[4.6, 3.6], [-0.2, 0.8]]), 8), empty(0, float)) @test def test_inv(a, expected): assert (round(alg.inv(a), 2) == round(expected, 2)).all() test_inv(np.array([[1, 2], [3, 4]]), np.array([[-2., 1.], [1.5, -0.5]])) test_inv(np.array([[[1., 2.], [3., 4.]], [[1., 3.], [3., 5.]]]), np.array([[[-2., 1.], [1.5, -0.5]], [[-1.25, 0.75], [0.75, -0.25]]])) @test def test_pinv(a, expected, hermitian=False): if not hermitian: assert np.allclose(a, np.dot(a, np.dot(alg.pinv(a), a))) else: assert np.allclose(alg.pinv(a, hermitian=hermitian)[0][0], expected) test_pinv( np.array([[ -0.11701212, 2.00785233, -0.18970224, -1.21453479, 0.65005008, 1.54383399 ], [ -0.77042552, -2.61734354, -1.94324215, 1.21451409, 0.03406445, -0.02870602 ], [ 0.77946711, 1.36558027, 1.41045231, -0.85389799, -0.61840915, 0.26268489 ]]), 0) test_pinv(np.array([[1, 2], [3, 4], [5, 6]]), 0) test_pinv(np.array([[1, 2], [3, 4]]), -0.79999999, True) @test def test_tensorinv(): a = np.eye(4 * 6) x = a.reshape((4, 6, 8, 3)) ainv = alg.tensorinv(x) b = np.array([[ 1.06690372, -1.18093448, 0.59943811, -0.21417083, -0.74814399, 0.22962492 ], [ -0.32909058, -0.06753367, 0.67559711, -1.13889073, -0.24667704, 0.19231421 ], [ 0.07065959, 0.48421004, -0.24765094, 0.15383651, 0.29206018, 0.24088399 ], [ -0.46943209, -0.95985002, -0.38301866, 0.01950948, -1.32758378, 0.84240223 ]]) d = (np.tensordot(ainv, b)).reshape((8, 3)) assert np.allclose(d, alg.tensorsolve(x, b)) test_tensorinv() @test def test_linalg_non_c_contig(x1, x2, expected): assert (np.matmul(x1, x2)[0] == expected[0]).all() fcontig_matrix1 = (np.array([[1., 3., 6.], [2., 4., 9.]])).T fcontig_matrix2 = (np.array([[1.4, 8.8], [3.6, 2.21]])).T test_linalg_non_c_contig( fcontig_matrix1, fcontig_matrix2, np.array([[19.0, 8.02], [39.4, 19.64], [87.6, 41.49]])) regular_matrix = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) noncontig_matrix1 = regular_matrix[::2, ::2] noncontig_matrix2 = regular_matrix[:2, 1::2] test_linalg_non_c_contig(noncontig_matrix1, noncontig_matrix2, np.array([[20, 28], [84, 124]], dtype=int)) @test def test_matmul_contigs(dtype: type): a = np.array([[1, 2, -3], [4, -5, 6]], dtype) b = np.array([[1, 2, 3, 4], [4, 5, 6, 7], [9, 8, 7, 6]], dtype) x = np.array([[1, 2], [-3, 4], [-5, 6]], dtype).T y = np.array([[1, 2, 3], [4, 4, 5], [6, 7, 9], [8, 7, 6]], dtype).T a1 = np.array([[1, 0, 2, 0, -3], [0] * 5, [4, 0, -5, 0, 6]], dtype)[::2, ::2] b1 = np.array([[1, 0, 2, 0, 3, 0, 4], [0] * 7, [4, 0, 5, 0, 6, 0, 7], [0] * 7, [9, 0, 8, 0, 7, 0, 6]], dtype)[::2, ::2] x1 = np.array([[1, 0, 2], [0] * 3, [-3, 0, 4], [0] * 3, [-5, 0, 6]], dtype).T[::2, ::2] y1 = np.array([[1, 0, 2, 0, 3], [0] * 5, [4, 0, 4, 0, 5], [0] * 5, [6, 0, 7, 0, 9], [0] * 5, [8, 0, 7, 0, 6]], dtype).T[::2, ::2] ab = np.array([[-18, -12, -6, 0], [38, 31, 24, 17]], dtype) xy = np.array([[-20, -33, -60, -43], [28, 54, 94, 80]], dtype) ay = np.array([[-4, -3, -7, 4], [12, 26, 43, 33]], dtype) xb = np.array([[-56, -53, -50, -47], [72, 72, 72, 72]], dtype) oc = np.empty((2, 4), order='C', dtype=dtype) of = np.empty((2, 4), order='F', dtype=dtype) assert np.array_equal(a @ b, ab) assert np.array_equal(x @ y, xy) assert np.array_equal(a @ y, ay) assert np.array_equal(x @ b, xb) assert np.array_equal(np.matmul(a, b, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(x, y, out=of), xy) assert np.array_equal(of, xy) assert np.array_equal(np.matmul(a, y, out=oc), ay) assert np.array_equal(oc, ay) assert np.array_equal(np.matmul(x, b, out=oc), xb) assert np.array_equal(oc, xb) assert np.array_equal(a1 @ b1, ab) assert np.array_equal(x1 @ y1, xy) assert np.array_equal(a1 @ y1, ay) assert np.array_equal(x1 @ b1, xb) assert np.array_equal(np.matmul(a1, b1, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(x1, y1, out=oc), xy) assert np.array_equal(oc, xy) assert np.array_equal(np.matmul(a1, y1, out=oc), ay) assert np.array_equal(oc, ay) assert np.array_equal(np.matmul(x1, b1, out=oc), xb) assert np.array_equal(oc, xb) @test def test_matmul_vec(dtype: type): # vec-mat a = np.array([11, 22, 33], dtype) b = np.array([[1, 2], [2, 3], [3, -5]], dtype) y = np.array([[1, 2, 2], [3, 3, -5]], dtype).T a1 = np.array([11, 0, 22, 0, 33], dtype)[::2] b1 = np.array([[1, 0, 2], [0] * 3, [2, 0, 3], [0] * 3, [3, 0, -5]], dtype)[::2, ::2] y1 = np.array([[1, 0, 2, 0, 2], [0] * 5, [3, 0, 3, 0, -5]], dtype).T[::2, ::2] ab = np.array([154, -77], dtype) ay = np.array([121, -66], dtype) oc = np.empty(2, order='C', dtype=dtype) of = np.empty(2, order='F', dtype=dtype) assert np.array_equal(a @ b, ab) assert np.array_equal(a @ y, ay) assert np.array_equal(a1 @ b1, ab) assert np.array_equal(a1 @ y1, ay) assert np.array_equal(a @ b1, ab) assert np.array_equal(a1 @ y, ay) assert np.array_equal(np.matmul(a, b, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(a, y, out=of), ay) assert np.array_equal(of, ay) assert np.array_equal(np.matmul(a1, b1, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(a1, y1, out=oc), ay) assert np.array_equal(oc, ay) assert np.array_equal(np.matmul(a, b1, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(a1, y, out=of), ay) assert np.array_equal(of, ay) # mat-vec a = np.array([[11, 22, 33], [-1, -2, -3]], dtype) x = np.array([[11, 22], [33, -1], [-2, -3]], dtype).T b = np.array([1, 2, 3], dtype) a1 = np.array([[11, 0, 22, 0, 33], [0] * 5, [-1, 0, -2, 0, -3]], dtype)[::2, ::2] x1 = np.array([[11, 0, 22], [0] * 3, [33, 0, -1], [0] * 3, [-2, 0, -3]], dtype).T[::2, ::2] b1 = np.array([1, 0, 2, 0, 3], dtype)[::2] ab = np.array([154, -14], dtype) xb = np.array([71, 11], dtype) oc = np.empty(2, order='C', dtype=dtype) of = np.empty(2, order='F', dtype=dtype) assert np.array_equal(a @ b, ab) assert np.array_equal(x @ b, xb) assert np.array_equal(a1 @ b1, ab) assert np.array_equal(x1 @ b1, xb) assert np.array_equal(a @ b1, ab) assert np.array_equal(x1 @ b, xb) assert np.array_equal(np.matmul(a, b, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(x, b, out=of), xb) assert np.array_equal(of, xb) assert np.array_equal(np.matmul(a1, b1, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(x1, b1, out=oc), xb) assert np.array_equal(oc, xb) assert np.array_equal(np.matmul(a1, b1, out=oc), ab) assert np.array_equal(oc, ab) assert np.array_equal(np.matmul(x1, b1, out=oc), xb) assert np.array_equal(oc, xb) @test def test_matmul_syrk(dtype: type): a = np.array([[1, 3, -2], [-1, 4, -5]], dtype=dtype) assert np.array_equal(a @ a.T, np.array([[14, 21], [21, 42]], dtype=dtype)) assert np.array_equal( a.T @ a, np.array([[2, -1, 3], [-1, 25, -26], [3, -26, 29]], dtype=dtype)) a = np.array([[1, 3, -2], [-1, 4, -5]], dtype=dtype).T assert np.array_equal( a @ a.T, np.array([[2, -1, 3], [-1, 25, -26], [3, -26, 29]], dtype=dtype)) assert np.array_equal(a.T @ a, np.array([[14, 21], [21, 42]], dtype=dtype)) a = np.array([[1, 3], [-2, -1], [4, -5]], dtype=dtype) assert np.array_equal( a @ a.T, np.array([[10, -5, -11], [-5, 5, -3], [-11, -3, 41]], dtype=dtype)) assert np.array_equal(a.T @ a, np.array([[21, -15], [-15, 35]], dtype=dtype)) a = np.array([[1, 3], [-2, -1], [4, -5]], dtype=dtype).T assert np.array_equal(a @ a.T, np.array([[21, -15], [-15, 35]], dtype=dtype)) assert np.array_equal( a.T @ a, np.array([[10, -5, -11], [-5, 5, -3], [-11, -3, 41]], dtype=dtype)) @test def test_dot_matmul_higher_dims(dtype: type): # MatMul A = np.arange(2*3*4, dtype=dtype).reshape(2, 3, 4) B = np.arange(2*3*4, dtype=dtype).reshape(2, 4, 3) C = A @ B assert np.array_equal(C, np.array([[[42.0, 48.0, 54.0], [114.0, 136.0, 158.0], [186.0, 224.0, 262.0]], [[906.0, 960.0, 1014.0], [1170.0, 1240.0, 1310.0], [1434.0, 1520.0, 1606.0]]], dtype)) D = np.zeros((2, 3, 3), dtype) E = np.matmul(A, B, out=D) assert np.array_equal(C, D) assert D.data == E.data A = np.arange(3*4, dtype=dtype).reshape(3, 4) B = np.arange(2*3*4, dtype=dtype).reshape(2, 4, 3) C = A @ B assert np.array_equal(C, np.array([[[42.0, 48.0, 54.0], [114.0, 136.0, 158.0], [186.0, 224.0, 262.0]], [[114.0, 120.0, 126.0], [378.0, 400.0, 422.0], [642.0, 680.0, 718.0]]], dtype)) D = np.zeros((2, 3, 3), dtype) E = np.matmul(A, B, out=D) assert np.array_equal(C, D) assert D.data == E.data A = np.arange(2*3*4, dtype=dtype).reshape(2, 3, 4) B = np.arange(3*4, dtype=dtype).reshape(4, 3) C = A @ B assert np.array_equal(C, np.array([[[42.0, 48.0, 54.0], [114.0, 136.0, 158.0], [186.0, 224.0, 262.0]], [[258.0, 312.0, 366.0], [330.0, 400.0, 470.0], [402.0, 488.0, 574.0]]], dtype)) D = np.zeros((2, 3, 3), dtype) E = np.matmul(A, B, out=D) assert np.array_equal(C, D) assert D.data == E.data # Dot A = np.arange(2*3*4, dtype=dtype).reshape(2, 3, 4) B = np.arange(2*3*4, dtype=dtype).reshape(2, 4, 3) C = np.dot(A, B) assert np.array_equal(C, np.array([[[[42.0, 48.0, 54.0], [114.0, 120.0, 126.0]], [[114.0, 136.0, 158.0], [378.0, 400.0, 422.0]], [[186.0, 224.0, 262.0], [642.0, 680.0, 718.0]]], [[[258.0, 312.0, 366.0], [906.0, 960.0, 1014.0]], [[330.0, 400.0, 470.0], [1170.0, 1240.0, 1310.0]], [[402.0, 488.0, 574.0], [1434.0, 1520.0, 1606.0]]]], dtype)) D = np.zeros((2, 3, 2, 3), dtype) E = np.dot(A, B, out=D) assert np.array_equal(C, D) assert D.data == E.data A = np.arange(3*4, dtype=dtype).reshape(3, 4) B = np.arange(2*3*4, dtype=dtype).reshape(2, 4, 3) C = np.dot(A, B) assert np.array_equal(C, np.array([[[42.0, 48.0, 54.0], [114.0, 120.0, 126.0]], [[114.0, 136.0, 158.0], [378.0, 400.0, 422.0]], [[186.0, 224.0, 262.0], [642.0, 680.0, 718.0]]], dtype)) D = np.zeros((3, 2, 3), dtype) E = np.dot(A, B, out=D) assert np.array_equal(C, D) assert D.data == E.data A = np.arange(2*3*4, dtype=dtype).reshape(2, 3, 4) B = np.arange(3*4, dtype=dtype).reshape(4, 3) C = np.dot(A, B) assert np.array_equal(C, np.array([[[42.0, 48.0, 54.0], [114.0, 136.0, 158.0], [186.0, 224.0, 262.0]], [[258.0, 312.0, 366.0], [330.0, 400.0, 470.0], [402.0, 488.0, 574.0]]], dtype)) D = np.zeros((2, 3, 3), dtype) E = np.matmul(A, B, out=D) assert np.array_equal(C, D) assert D.data == E.data test_matmul_contigs(np.float64) test_matmul_contigs(np.float32) test_matmul_contigs(np.complex64) test_matmul_contigs(np.complex128) test_matmul_contigs(np.int64) test_matmul_contigs(np.int32) test_matmul_vec(np.float64) test_matmul_vec(np.float32) test_matmul_vec(np.complex64) test_matmul_vec(np.complex128) test_matmul_vec(np.int64) test_matmul_vec(np.int32) test_matmul_syrk(np.float64) test_matmul_syrk(np.float32) test_matmul_syrk(np.complex128) test_matmul_syrk(np.complex64) test_matmul_syrk(np.int64) test_matmul_syrk(np.int32) test_dot_matmul_higher_dims(np.float64) test_dot_matmul_higher_dims(np.float32) test_dot_matmul_higher_dims(np.complex128) test_dot_matmul_higher_dims(np.complex64) test_dot_matmul_higher_dims(np.int64) test_dot_matmul_higher_dims(np.int32)