fork download
  1. import numpy as np
  2.  
  3. # Data setup
  4. A_data = np.array([
  5. 1, 2, 3, 4, 5, 6, 7, 8,
  6. 9, 10, 11, 12, 13, 14, 15, 16,
  7. 17, 18, 19, 20, 21, 22, 23, 24
  8. ], dtype=np.float32)
  9. B_data = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32)
  10.  
  11. # Shape in C++
  12. # Tensor<float> A({2, 3, 4}, ...)
  13. # Tensor<float> B({4, 2}, ...)
  14.  
  15. A = A_data.reshape(2, 3, 4, order='F') # Fortran order = column-major
  16. B = B_data.reshape(4, 2, order='F')
  17.  
  18. print("A shape:", A.shape)
  19. print("B shape:", B.shape)
  20. print("\nA (column-major order):")
  21. print(A)
  22. print("\nB (column-major order):")
  23. print(B)
  24.  
  25. # tensordot over axis=1: last axis of A, first axis of B
  26. result = np.tensordot(A, B, axes=1)
  27. print("\nResult shape:", result.shape)
  28.  
  29. # Output in column-major flattening
  30. result_flat = result.flatten(order='F')
  31. print("\nResult (flattened F):", result_flat)
  32.  
Success #stdin #stdout 0.08s 23744KB
stdin
Standard input is empty
stdout
('A shape:', (2, 3, 4))
('B shape:', (4, 2))

A (column-major order):
[[[ 1.  7. 13. 19.]
  [ 3.  9. 15. 21.]
  [ 5. 11. 17. 23.]]

 [[ 2.  8. 14. 20.]
  [ 4. 10. 16. 22.]
  [ 6. 12. 18. 24.]]]

B (column-major order):
[[1. 5.]
 [2. 6.]
 [3. 7.]
 [4. 8.]]
('\nResult shape:', (2, 3, 2))
('\nResult (flattened F):', array([130., 140., 150., 160., 170., 180., 290., 316., 342., 368., 394.,
       420.], dtype=float32))