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

Fix Metal API validation for empty concat (#1183)

Browse Source
This commit is contained in:
Awni Hannun 2024-06-04 13:17:08 -07:00 committed by GitHub
parent 375a8bbdcc
commit 83b11bc58d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/backend/metal/copy.cpp
View File

@ -33,9 +33,6 @@ void copy_gpu(const array& in, array& out, CopyType ctype, const Stream& s) {
} else { } else {
out.set_data(allocator::malloc_or_wait(out.nbytes())); out.set_data(allocator::malloc_or_wait(out.nbytes()));
} }
if (out.size() == 0) {
return;
}
if (ctype == CopyType::GeneralGeneral) { if (ctype == CopyType::GeneralGeneral) {
ctype = CopyType::General; ctype = CopyType::General;
} }
@ -57,6 +54,10 @@ void copy_gpu_inplace(
int64_t out_offset, int64_t out_offset,
CopyType ctype, CopyType ctype,
const Stream& s) { const Stream& s) {
if (out.size() == 0) {
return;
}
// Try to collapse contiguous dims // Try to collapse contiguous dims
auto [shape, strides] = collapse_contiguous_dims( auto [shape, strides] = collapse_contiguous_dims(
data_shape, std::vector{strides_in_pre, strides_out_pre}); data_shape, std::vector{strides_in_pre, strides_out_pre});

6
python/tests/test_ops.py
View File

@ -1533,6 +1533,12 @@ class TestOps(mlx_tests.MLXTestCase):
b = mx.array([1, 2]) b = mx.array([1, 2])
mx.concatenate([a, b], axis=0) mx.concatenate([a, b], axis=0)
# Cocnatenate with 0-sized array
a = mx.zeros((2, 0, 2))
b = mx.zeros((2, 2, 2))
out = mx.concatenate([a, b], axis=1)
self.assertTrue(mx.array_equal(out, b))
def test_meshgrid(self): def test_meshgrid(self):
x = mx.array([1, 2, 3], dtype=mx.int32) x = mx.array([1, 2, 3], dtype=mx.int32)
y = np.array([1, 2, 3], dtype=np.int32) y = np.array([1, 2, 3], dtype=np.int32)