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

Minor updates to address a few issues (#537)

Browse Source 
* docs on arg indices return type

* arange with nan

* undo isort
This commit is contained in:
Awni Hannun
2024-01-23 22:24:41 -08:00
committed by GitHub
parent 4fe2fa2a64
commit f30e63353a
3 changed files with 23 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/ops.cpp
View File

@@ -79,7 +79,14 @@ array arange(
msg << bool_ << " not supported for arange.";
throw std::invalid_argument(msg.str());
}
int size = std::max(static_cast<int>(std::ceil((stop - start) / step)), 0);
if (std::isnan(start) || std::isnan(step) || std::isnan(stop)) {
throw std::invalid_argument("[arange] Cannot compute length.");
}
double real_size = std::ceil((stop - start) / step);
if (std::isnan(real_size)) {
throw std::invalid_argument("[arange] Cannot compute length.");
}
int size = std::max(static_cast<int>(real_size), 0);
return array(
{size},
dtype,