Kỹ thuật Log-Sum-Exp tránh tràn số dấu phẩy động (Log-Sum-Exp Trick)

Mô tả bài toán

Trong học máy, khi tính chuẩn hóa xác suất hoặc hàm Cross-Entropy Loss, ta thường xuyên phải tính: log ( ∑_i=1^V e^z_i ) Nếu tính ngây thơ bằng np.log(np.sum(np.exp(z))):

• Nếu z_i > 709 (với float64) hoặc z_i > 88 (với float32), e^z_i sẽ tràn số thành inf, khiến kết quả trở thành inf.
• Nếu tất cả z_i < -709, e^z_i sẽ triệt tiêu dưới số 0 (underflow), khiến kết quả là log(0) = -inf.

Kỹ thuật Log-Sum-Exp (LSE) triệt tiêu hoàn toàn lỗi số học bằng cách trừ đi giá trị cực đại m = max_j z_j: LSE(z) = m + log ( ∑_i=1^V e^{z_i - m} )

Hãy viết hàm stable_logsumexp(Z: np.ndarray, axis: int = -1, keepdims: bool = False) -> np.ndarray:

• Nhận tensor Z bất kỳ số chiều.
• Tính toán Log-Sum-Exp dọc theo trục axis được chỉ định.
• Đảm bảo tính toán hoàn toàn ổn định trên kiểu dữ liệu np.float64.
• Trả về mảng NumPy float64.

Ràng buộc

• Thời gian chạy tối đa: 6000ms.
• 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ả.