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

support arange for bfloat16 (#245)

Browse Source
This commit is contained in:
Daniel Strobusch
2023-12-21 23:33:43 +01:00
committed by GitHub
parent 2c7df6795e
commit 794feb83df
2 changed files with 17 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/metal/primitives.cpp
View File

@@ -215,7 +215,8 @@ void Arange::eval_gpu(const std::vector<array>& inputs, array& out) {
arange_set_scalars<float>(start_, start_ + step_, compute_encoder);
break;
case bfloat16:
throw std::runtime_error("[Arange::eval_gpu] Does not support bfloat16");
arange_set_scalars<bfloat16_t>(start_, start_ + step_, compute_encoder);
break;
case complex64:
throw std::runtime_error("[Arange::eval_gpu] Does not support complex64");
}