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 Packages Projects Releases Wiki Activity

Address comments

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-07-07 17:03:28 -07:00
parent 8ea5729ee4
commit 1c589298ec
2 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
mlx/backend/metal/indexing.cpp
View File

@@ -1,6 +1,5 @@
// Copyright © 2023-2024 Apple Inc. // Copyright © 2023-2024 Apple Inc.
#include <fmt/format.h> #include <fmt/format.h>
#include <iostream>
#include "mlx/backend/common/compiled.h" #include "mlx/backend/common/compiled.h"
#include "mlx/backend/common/utils.h" #include "mlx/backend/common/utils.h"

5
mlx/ops.cpp
View File

@@ -4676,6 +4676,11 @@ array segmented_mm(
throw std::invalid_argument(msg.str()); throw std::invalid_argument(msg.str());
} }
if (!issubdtype(segments.dtype(), integer)) {
throw std::invalid_argument(
"[segmented_mm] Got segments with invalid dtype. Segments must be integral.");
}
a = astype(a, out_type, s); a = astype(a, out_type, s);
b = astype(b, out_type, s); b = astype(b, out_type, s);
segments = astype(segments, uint32, s); segments = astype(segments, uint32, s);