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

Explicit barriers with concurrent dispatch (#977)

Browse Source
This commit is contained in:
Awni Hannun
2024-04-10 21:45:31 -07:00
committed by GitHub
parent 8580d997ff
commit 12d4507ee3
21 changed files with 326 additions and 267 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
mlx/ops.cpp
View File

@@ -1050,19 +1050,30 @@ array transpose(
for (auto& ax : axes) {
ax = ax < 0 ? ax + a.ndim() : ax;
}
std::set dims(axes.begin(), axes.end());
if (dims.size() != axes.size()) {
throw std::invalid_argument("Repeat axes not allowed in transpose.");
if (axes.size() != a.ndim()) {
std::ostringstream msg;
msg << "[transpose] Recived " << axes.size() << " axes for array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
if (dims.size() != a.ndim() ||
a.ndim() > 0 &&
(*dims.begin() != 0 || *dims.rbegin() != (a.ndim() - 1))) {
throw std::invalid_argument("Transpose axes don't match array dimensions.");
// Check in bounds and for duplicates
std::vector<int> shape(axes.size(), 0);
for (auto& ax : axes) {
if (ax < 0 || ax >= a.ndim()) {
std::ostringstream msg;
msg << "[transpose] Invalid axis (" << ax << ") for array with "
<< a.ndim() << " dimensions.";
throw std::invalid_argument(msg.str());
}
if (shape[ax] != 0) {
throw std::invalid_argument("[transpose] Repeat axes not allowed.");
}
shape[ax] = 1;
}
std::vector<int> shape;
shape.reserve(axes.size());
for (auto ax : axes) {
shape.push_back(a.shape()[ax]);
for (int i = 0; i < axes.size(); ++i) {
shape[i] = a.shape()[axes[i]];
}
return array(
std::move(shape),