Thao Tác Tensor với Ký Hiệu Einstein (torch.einsum) trong LLM

Mô tả bài toán

Trong mã nguồn của các kiến trúc LLM hiện đại (như LLaMA, GPT, Mistral), torch.einsum là công cụ ngắn gọn và mạnh mẽ nhất để biểu diễn các phép co tensor đa chiều (tensor contraction), tích vô hướng qua từng đầu chú ý (Multi-Head Attention scores), và tính vết ma trận (trace).

Yêu cầu

Viết hàm execute_einsum_operation(operation_type: str, *tensors: np.ndarray) -> np.ndarray:

• "batch_attention_scores": Nhận 2 tensor Q shape (B, H, Q_len, D) và K shape (B, H, K_len, D). Tính ma trận điểm chú ý raw scores qua phép nhân Q K^T:

Shape đầu ra: (B, H, Q_len, K_len).

• "batched_diagonal": Nhận 1 tensor ma trận theo lô M shape (B, N, N). Trích xuất các đường chéo chính:

Shape đầu ra: (B, N).

• "matrix_trace": Nhận 1 ma trận vuông A shape (N, N). Tính vết (tổng các phần tử đường chéo chính):

Kết quả trả về là một số thực vô hướng kiểu float64.

• Nếu operation_type không hợp lệ, raise ValueError.

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.