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

Fix four step fft

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-05-08 14:14:59 -07:00
parent 2a41caa00e
commit 83762691ba
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/backend/metal/fft.cpp
View File

@@ -1117,8 +1117,11 @@ void fft_four_step_inplace(
} }
// Get the kernel // Get the kernel
auto in_type = in.dtype() == float32 ? "float" : "float2"; auto to_type = [](const array& x) {
auto out_type = out.dtype() == float32 ? "float" : "float2"; return x.dtype() == float32 ? "float" : "float2";
};
auto in_type = step == 0 ? to_type(in) : to_type(intermediate);
auto out_type = step == 0 ? to_type(intermediate) : to_type(out);
std::string hash_name; std::string hash_name;
std::string kname; std::string kname;
kname.reserve(64); kname.reserve(64);