Bentuk hasil perkalian matriks

Kembalikan bentuk hasil A @ B (seperti np.matmul) dari bentuk a dan b, atau None jika tidak kompatibel. Kedua bentuk selalu punya panjang minimal 2.

[2, 3]    @ [3, 4] → [2, 4]
[5, 2, 3] @ [3, 4] → [5, 2, 4]   // batch 5, matriks yang sama dipakai untuk semuanya
[2, 3]    @ [4, 3] → None        // 3 ≠ 4

Matmul berbatch

Di transformer, attention menghitung Q @ K.T untuk setiap batch dan setiap head sekaligus: (B, H, T, d) @ (B, H, d, T) → (B, H, T, T). NumPy (dan PyTorch) memperlakukan dua dimensi terakhir sebagai matriks, dan sisanya sebagai “tumpukan” matriks yang di-broadcast.

Cara mengecek

Dimensi dalam: a[-1] harus sama dengan b[-2].
Hasil: batch + [a[-2], b[-1]], dengan batch = broadcast(a[:-2], b[:-2]).

Biasakan menulis komentar bentuk di setiap baris, misalnya # (B, T, d) @ (d, 3d) → (B, T, 3d). Kebiasaan kecil ini menghemat berjam-jam debugging.

Penjelasan

Yang diminta. Bentuk hasil A @ B untuk bentuk berdimensi ≥ 2, atau None jika tidak bisa dikalikan.

Ide kunci. @ memperlakukan dua dimensi terakhir sebagai matriks, dan semua dimensi di depannya sebagai “tumpukan” (batch). Jadi soal ini adalah dua soal kecil:

Bagian matriks: (n, k) @ (k, m) → (n, m). Syaratnya a[-1] == b[-2] (dimensi dalam sama).
Bagian batch: a[:-2] dan b[:-2] di-broadcast dengan aturan dari soal sebelumnya.

Langkah di kode:

Cek dimensi dalam dulu, karena paling murah: if a[-1] != b[-2]: return None.
batch = broadcast(a[:-2], b[:-2]) memakai ulang fungsi broadcasting. Untuk [5, 2, 3] @ [3, 4], batch-nya broadcast([5], []) = [5].
Jika batch tidak kompatibel, None.
Gabungkan: batch + [a[-2], b[-1]], yaitu baris dari A dan kolom dari B.

Yang dikembalikan: list bentuk, misalnya [4, 5, 2, 6], atau None.

Kesalahan umum: mengecek a[-1] == b[-1] (itu syarat untuk A @ B.T); lupa bahwa batch juga boleh di-broadcast ([4, 1, ...] dengan [5, ...]).