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

Change structured binding capture

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-06-29 04:45:25 -07:00
parent aa7841814c
commit 4aca09339d
2 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/backend/cuda/binary_two.cu
View File

@@ -153,7 +153,10 @@ void binary_op_gpu_inplace(
out_a.data_size() > INT32_MAX, out_a.data_size() > INT32_MAX,
[&](auto large) { [&](auto large) {
using IdxT = std::conditional_t<large(), int64_t, int32_t>; using IdxT = std::conditional_t<large(), int64_t, int32_t>;
auto [shape, strides] = collapse_contiguous_dims(a, b, out_a); Shape shape;
std::vector<Strides> strides;
std::tie(shape, strides) =
collapse_contiguous_dims(a, b, out_a);
auto& a_strides = strides[0]; auto& a_strides = strides[0];
auto& b_strides = strides[1]; auto& b_strides = strides[1];
int ndim = shape.size(); int ndim = shape.size();

4
mlx/backend/cuda/ternary.cu
View File

@@ -102,7 +102,9 @@ void ternary_op_gpu_inplace(
c.data_size() > INT32_MAX || out.data_size() > INT32_MAX, c.data_size() > INT32_MAX || out.data_size() > INT32_MAX,
[&](auto large) { [&](auto large) {
using IdxT = std::conditional_t<large(), int64_t, int32_t>; using IdxT = std::conditional_t<large(), int64_t, int32_t>;
auto [shape, strides] = collapse_contiguous_dims(a, b, c, out); Shape shape;
std::vector<Strides> strides;
std::tie(shape, strides) = collapse_contiguous_dims(a, b, c, out);
auto& a_strides = strides[0]; auto& a_strides = strides[0];
auto& b_strides = strides[1]; auto& b_strides = strides[1];
auto& c_strides = strides[2]; auto& c_strides = strides[2];