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

1import numpy as np 

2 

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 

8 

9# make "random" input arrays 

10 

11 

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. 

15 

16 # built-in dot product 

17 AB1 = np.dot(A, B) 

18 

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]) 

26 

27 # test for consistency 

28 assert np.all(AB1 == AB2)