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

fix: handle linspace function when num is 1 (#602)

Browse Source 
* fix: handle linspace function when num is 1

* add comment

* fix test case

* remove breakpoint
This commit is contained in:
Avikant Srivastava
2024-02-05 00:33:49 +05:30
committed by GitHub
parent 4fd2fb84a6
commit 11a9fd40f0
2 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/ops.cpp
View File

@@ -148,6 +148,9 @@ array linspace(
msg << "[linspace] number of samples, " << num << ", must be non-negative.";
throw std::invalid_argument(msg.str());
}
if (num == 1) {
return astype(array({start}), dtype, to_stream(s));
}
array sequence = arange(0, num, float32, to_stream(s));
float step = (stop - start) / (num - 1);
return astype(