zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-24 02:41:19 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

More lenient mask type check in SDPA (#1723)

Browse Source 
* check mask type

* require promotion
This commit is contained in:
Alex Barron 2024-12-18 19:41:38 -08:00 committed by GitHub
parent ed4ec81bca
commit f17536af9c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/fast.cpp
View File

@ -586,9 +586,9 @@ array scaled_dot_product_attention(
throw std::invalid_argument(msg.str()); throw std::invalid_argument(msg.str());
} }
if (mask && (*mask).dtype() != final_type) { if (mask && promote_types((*mask).dtype(), final_type) != final_type) {
std::ostringstream msg; std::ostringstream msg;
msg << "[scaled_dot_product_attention] Mask should match output type. " msg << "[scaled_dot_product_attention] Mask type must promote to output type. "
<< final_type << "."; << final_type << ".";
throw std::invalid_argument(msg.str()); throw std::invalid_argument(msg.str());
} }