fix rebase

ml-explore · Sep 26, 2024 · 15ca8cd · 15ca8cd
1 parent 2dbda67
commit 15ca8cd
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 13 deletions.
diff --git a/mlx/backend/common/utils.cpp b/mlx/backend/common/utils.cpp
@@ -86,8 +86,7 @@ collapse_contiguous_dims(
 }
 
 template <typename StrideT>
-std::tuple<std::vector<int>, std::vector<StrideT>>
-collapse_contiguous_dims_impl(
+std::pair<std::vector<int>, std::vector<StrideT>> collapse_contiguous_dims_impl(
     const std::vector<int>& shape,
     const std::vector<StrideT>& strides,
     StrideT size_cap) {
@@ -112,24 +111,24 @@ collapse_contiguous_dims_impl(
     }
   }
 
-  return std::make_tuple(collapsed_shape, collapsed_strides);
+  return std::make_pair(collapsed_shape, collapsed_strides);
 }
 
-std::tuple<std::vector<int>, std::vector<int64_t>> collapse_contiguous_dims(
+std::pair<std::vector<int>, std::vector<int64_t>> collapse_contiguous_dims(
     const std::vector<int>& shape,
     const std::vector<int64_t>& strides,
     int64_t size_cap /* = std::numeric_limits<int32_t>::max() */) {
   return collapse_contiguous_dims_impl<int64_t>(shape, strides, size_cap);
 }
 
-std::tuple<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
+std::pair<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
     const std::vector<int>& shape,
     const std::vector<size_t>& strides,
     size_t size_cap /* = std::numeric_limits<int32_t>::max() */) {
   return collapse_contiguous_dims_impl<size_t>(shape, strides, size_cap);
 }
 
-std::tuple<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
+std::pair<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
     const array& a,
     size_t size_cap /* = std::numeric_limits<int32_t>::max()*/) {
   return collapse_contiguous_dims_impl<size_t>(

diff --git a/mlx/backend/common/utils.h b/mlx/backend/common/utils.h
@@ -73,15 +73,15 @@ inline auto collapse_contiguous_dims(Arrays&&... xs) {
 }
 
 // The single array version of the above.
-std::tuple<std::vector<int>, std::vector<int64_t>> collapse_contiguous_dims(
+std::pair<std::vector<int>, std::vector<int64_t>> collapse_contiguous_dims(
     const std::vector<int>& shape,
     const std::vector<int64_t>& strides,
     int64_t size_cap = std::numeric_limits<int32_t>::max());
-std::tuple<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
+std::pair<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
     const std::vector<int>& shape,
     const std::vector<size_t>& strides,
     size_t size_cap = std::numeric_limits<int32_t>::max());
-std::tuple<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
+std::pair<std::vector<int>, std::vector<size_t>> collapse_contiguous_dims(
     const array& a,
     size_t size_cap = std::numeric_limits<int32_t>::max());
 

diff --git a/mlx/backend/metal/unary.cpp b/mlx/backend/metal/unary.cpp
@@ -28,10 +28,7 @@ void unary_op_gpu_inplace(
 
   auto maybe_collapse = [contig, &in, &out]() {
     if (!contig) {
-      auto [shape, strides] = collapse_contiguous_dims(
-          {in, out},
-          /* size_cap = */ INT32_MAX);
-      return std::make_pair(shape, strides[0]);
+      return collapse_contiguous_dims(in);
     } else {
       return std::make_pair(std::vector<int>{}, std::vector<size_t>{});
     }