fork download
  1. import numpy as np
  2.  
  3. # Vector dot product: (3,) @ (3,) -> scalar
  4. v1 = np.array([1.0, 2.0, 3.0], dtype=np.float32)
  5. v2 = np.array([4.0, 5.0, 6.0], dtype=np.float32)
  6. vec_dot_result = np.dot(v1, v2)
  7. print("Vector dot product result:", vec_dot_result) # scalar
  8.  
  9. print("---")
  10.  
  11. # Matrix multiply: (2,3) @ (3,2) -> (2,2)
  12. A = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32).reshape(2, 3, order='F')
  13. B = np.array([7, 8, 9, 10, 11, 12], dtype=np.float32).reshape(3, 2, order='F')
  14. mat_mul_result = np.matmul(A, B)
  15. print("Shape:", mat_mul_result.shape)
  16. print(mat_mul_result.flatten(order='F')) # column major flat
  17.  
  18. print("---")
  19.  
  20. # Vector-matrix multiply: (3,) @ (3,2) -> (2,)
  21. C = np.array([4, 5, 6], dtype=np.float32)
  22. D = np.array([7, 8, 9, 10, 11, 12], dtype=np.float32).reshape(3, 2, order='F')
  23. mat_mul_result2 = np.matmul(C, D)
  24. print("Shape:", mat_mul_result2.shape)
  25. print(mat_mul_result2.flatten(order='F'))
  26.  
  27. print("---")
  28.  
  29. # Matrix-vector multiply: (4,5) @ (5,) -> (4,)
  30. E = np.array(range(1, 21), dtype=np.float32).reshape(4, 5, order='F')
  31. F = np.array([2, 3, 4, 5, 6], dtype=np.float32)
  32. mat_mul_result3 = np.matmul(E, F)
  33. print("Shape:", mat_mul_result3.shape)
  34. print(mat_mul_result3.flatten(order='F'))
  35.  
  36. print("---")
  37.  
  38. # Batched matmul: (2,2,3) @ (2,3,2) -> (2,2,2)
  39. a = np.array([1,2,3,4,5,6,7,8,9,10,11,12], dtype=np.float32).reshape(2, 2, 3, order='F')
  40. b = np.array([1,2,3,4,5,6,7,8,9,10,11,12], dtype=np.float32).reshape(2, 3, 2, order='F')
  41. mat_mul_result4 = np.matmul(a, b)
  42. print("Shape:", mat_mul_result4.shape)
  43. print(mat_mul_result4.flatten(order='F'))
  44.  
  45. print("---")
  46.  
  47. # Broadcasting example: (3,1,2,3) @ (1,4,3,2) -> (3,4,2,2)
  48. c = np.array(list(range(1, 19)), dtype=np.float32).reshape(3, 1, 2, 3, order='F')
  49. d = np.array(list(range(1, 25)), dtype=np.float32).reshape(1, 4, 3, 2, order='F')
  50. mat_mul_result5 = np.matmul(c, d)
  51. print("Shape:", mat_mul_result5.shape)
  52. print(mat_mul_result5.flatten(order='F'))
  53.  
  54. print("---")
  55.  
  56. # Vector @ ND: (3,) @ (2,3,4) -> (2,4)
  57. vec = np.array([1, 2, 3], dtype=np.float32)
  58. mat = np.array(list(range(1, 25)), dtype=np.float32).reshape(2, 3, 4, order='F')
  59. mat_mul_result6 = np.matmul(vec, mat)
  60. print("Shape:", mat_mul_result6.shape)
  61. print(mat_mul_result6.flatten(order='F'))
  62.  
  63. print("---")
  64.  
  65. # ND @ Vector: (2,3,4) @ (4,) -> (2,3)
  66. vec2 = np.array([1, 2, 3, 4], dtype=np.float32)
  67. mat2 = np.array(list(range(1, 25)), dtype=np.float32).reshape(2, 3, 4, order='F')
  68. mat_mul_result7 = np.matmul(mat2, vec2)
  69. print("Shape:", mat_mul_result7.shape)
  70. print(mat_mul_result7.flatten(order='F'))
  71.  
Success #stdin #stdout 0.07s 23720KB
stdin
Standard input is empty
stdout
('Vector dot product result:', 32.0)
---
('Shape:', (2, 2))
[ 76. 100. 103. 136.]
---
('Shape:', (2,))
[122. 167.]
---
('Shape:', (4,))
[220. 240. 260. 280.]
---
('Shape:', (2, 2, 2))
[ 61.  88.  79. 112. 151. 196. 205. 256.]
---
('Shape:', (3, 4, 2, 2))
[153. 168. 183. 174. 192. 210. 195. 216. 237. 216. 240. 264. 198. 213.
 228. 228. 246. 264. 258. 279. 300. 288. 312. 336. 405. 456. 507. 426.
 480. 534. 447. 504. 561. 468. 528. 588. 558. 609. 660. 588. 642. 696.
 618. 675. 732. 648. 708. 768.]
---
('Shape:', (2, 4))
[ 22.  28.  58.  64.  94. 100. 130. 136.]
---
('Shape:', (2, 3))
[130. 140. 150. 160. 170. 180.]