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

Fix arange with inf step (#686)

Browse Source 
* Fix case for step=inf in arange and add inf check for start/stop

* Add test cases for arange

* Update ops.cpp to include climits header

* Fix arange

* Fix formatting

* Refactor

* Add missing include
This commit is contained in:
Noah Farr
2024-02-23 15:18:15 +01:00
committed by GitHub
parent 126c9869c8
commit d729a1991b
2 changed files with 43 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
mlx/ops.cpp
View File

@@ -1,6 +1,7 @@
// Copyright © 2023-2024 Apple Inc.
#include <algorithm>
#include <climits>
#include <cmath>
#include <numeric>
#include <set>
@@ -73,10 +74,24 @@ array arange(
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)) {
if (std::isinf(start) || std::isinf(stop)) {
throw std::invalid_argument("[arange] Cannot compute length.");
}
// Check if start and stop specify a valid range because if not, we have to
// return an empty array
if (std::isinf(step) &&
(step > 0 && start < stop || step < 0 && start > stop)) {
return array({start}, dtype);
}
double real_size = std::ceil((stop - start) / step);
if (real_size > INT_MAX) {
throw std::invalid_argument("[arange] Maximum size exceeded.");
}
int size = std::max(static_cast<int>(real_size), 0);
return array(
{size},