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

check mask type (#1721)

Browse Source
This commit is contained in:
Alex Barron 2024-12-18 14:25:18 -08:00 committed by GitHub
parent 5548fcc96d
commit 49c34c4161
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/fast.cpp
View File

@ -586,6 +586,13 @@ array scaled_dot_product_attention(
throw std::invalid_argument(msg.str());
}
if (mask && (*mask).dtype() != final_type) {
std::ostringstream msg;
msg << "[scaled_dot_product_attention] Mask should match output type. "
<< final_type << ".";
throw std::invalid_argument(msg.str());
}
auto q = astype(queries, final_type, s);
auto k = astype(keys, final_type, s);
auto v = astype(values, final_type, s);