import numpy as np
# Data setup
A_data = np.array([
1, 2, 3, 4, 5, 6, 7, 8,
9, 10, 11, 12, 13, 14, 15, 16,
17, 18, 19, 20, 21, 22, 23, 24
], dtype=np.float32)
B_data = np.array([1, 2, 3, 4, 5, 6, 7, 8], dtype=np.float32)
# Shape in C++
# Tensor<float> A({2, 3, 4}, ...)
# Tensor<float> B({4, 2}, ...)
A = A_data.reshape(2, 3, 4, order='F') # Fortran order = column-major
B = B_data.reshape(4, 2, order='F')
print("A shape:", A.shape)
print("B shape:", B.shape)
print("\nA (column-major order):")
print(A)
print("\nB (column-major order):")
print(B)
# tensordot over axis=1: last axis of A, first axis of B
result = np.tensordot(A, B, axes=1)
print("\nResult shape:", result.shape)
# Output in column-major flattening
result_flat = result.flatten(order='F')
print("\nResult (flattened F):", result_flat)
aW1wb3J0IG51bXB5IGFzIG5wCgojIERhdGEgc2V0dXAKQV9kYXRhID0gbnAuYXJyYXkoWwogICAgMSwgMiwgMywgNCwgNSwgNiwgNywgOCwKICAgIDksIDEwLCAxMSwgMTIsIDEzLCAxNCwgMTUsIDE2LAogICAgMTcsIDE4LCAxOSwgMjAsIDIxLCAyMiwgMjMsIDI0Cl0sIGR0eXBlPW5wLmZsb2F0MzIpCkJfZGF0YSA9IG5wLmFycmF5KFsxLCAyLCAzLCA0LCA1LCA2LCA3LCA4XSwgZHR5cGU9bnAuZmxvYXQzMikKCiMgU2hhcGUgaW4gQysrCiMgVGVuc29yPGZsb2F0PiBBKHsyLCAzLCA0fSwgLi4uKQojIFRlbnNvcjxmbG9hdD4gQih7NCwgMn0sIC4uLikKCkEgPSBBX2RhdGEucmVzaGFwZSgyLCAzLCA0LCBvcmRlcj0nRicpICAjIEZvcnRyYW4gb3JkZXIgPSBjb2x1bW4tbWFqb3IKQiA9IEJfZGF0YS5yZXNoYXBlKDQsIDIsIG9yZGVyPSdGJykKCnByaW50KCJBIHNoYXBlOiIsIEEuc2hhcGUpCnByaW50KCJCIHNoYXBlOiIsIEIuc2hhcGUpCnByaW50KCJcbkEgKGNvbHVtbi1tYWpvciBvcmRlcik6IikKcHJpbnQoQSkKcHJpbnQoIlxuQiAoY29sdW1uLW1ham9yIG9yZGVyKToiKQpwcmludChCKQoKIyB0ZW5zb3Jkb3Qgb3ZlciBheGlzPTE6IGxhc3QgYXhpcyBvZiBBLCBmaXJzdCBheGlzIG9mIEIKcmVzdWx0ID0gbnAudGVuc29yZG90KEEsIEIsIGF4ZXM9MSkKcHJpbnQoIlxuUmVzdWx0IHNoYXBlOiIsIHJlc3VsdC5zaGFwZSkKCiMgT3V0cHV0IGluIGNvbHVtbi1tYWpvciBmbGF0dGVuaW5nCnJlc3VsdF9mbGF0ID0gcmVzdWx0LmZsYXR0ZW4ob3JkZXI9J0YnKQpwcmludCgiXG5SZXN1bHQgKGZsYXR0ZW5lZCBGKToiLCByZXN1bHRfZmxhdCkK