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

Array api (#1289)

Browse Source 
* some updates for numpy 2.0 and array api

* some updates for numpy 2.0 and array api

* fix array api doc
This commit is contained in:
Awni Hannun
2024-07-26 10:40:49 -07:00
committed by GitHub
parent e9e53856d2
commit 7b456fd2c0
6 changed files with 70 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/ops.cpp
View File

@@ -1839,15 +1839,6 @@ array argsort(const array& a, int axis, StreamOrDevice s /* = {} */) {
throw std::invalid_argument(msg.str());
}
// TODO: Fix GPU kernel
if (a.shape(axis) >= (1u << 21) && to_stream(s).device.type == Device::gpu) {
std::ostringstream msg;
msg << "[argsort] GPU sort cannot handle sort axis of >= 2M elements,"
<< " got array with sort axis size " << a.shape(axis) << "."
<< " Please place this operation on the CPU instead.";
throw std::runtime_error(msg.str());
}
return array(
a.shape(), uint32, std::make_shared<ArgSort>(to_stream(s), axis), {a});
}