Coverage for gpaw/test/linalg/test_dot.py: 100%
12 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1import numpy as np
3# Test that numpy.dot behaves as expected, i.e.
4# [A . B]_ijpq = sum_k A_ijk * B_pkq
5# Product sum is over last dimension of A and second-to-last dimension of B
6#
7# See numpy.tensordot for ultimate flexibility in choosing the pruduct-sum axes
9# make "random" input arrays
12def test_linalg_dot():
13 A = np.arange(6 * 2 * 5).reshape(6, 2, 5) - 3.
14 B = np.arange(3 * 5 * 4).reshape(3, 5, 4) + 5.
16 # built-in dot product
17 AB1 = np.dot(A, B)
19 # manual dot product
20 AB2 = np.empty(A.shape[:-1] + B.shape[:-2] + (B.shape[-1],))
21 for i in range(AB2.shape[0]):
22 for j in range(AB2.shape[1]):
23 for p in range(AB2.shape[2]):
24 for q in range(AB2.shape[3]):
25 AB2[i, j, p, q] = np.sum(A[i, j, :] * B[p, :, q])
27 # test for consistency
28 assert np.all(AB1 == AB2)