Gather mm new kernel and small refactoring (#2040)

2025-12-16 01:49:05 +08:00 · 2025-04-14 16:37:36 -07:00
parent e9e268336b
commit 99eefd2ec0
23 changed files with 1260 additions and 378 deletions
--- a/mlx/ops.h
+++ b/mlx/ops.h
@@ -1399,6 +1399,7 @@ array gather_mm(
    array b,
    std::optional<array> lhs_indices = std::nullopt,
    std::optional<array> rhs_indices = std::nullopt,
+    bool sorted_indices = false,
    StreamOrDevice s = {});

 /** Extract a diagonal or construct a diagonal array */