zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix bug with multiple attributes (#1348)

Browse Source 
Co-authored-by: Alex Barron <abarron22@apple.com>
This commit is contained in:
Alex Barron
2024-08-23 10:06:15 -07:00
committed by GitHub
parent 98b6ce3460
commit da8deb2b62
2 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_fast.py
View File

@@ -618,12 +618,12 @@ class TestFast(mlx_tests.MLXTestCase):
uint elem = thread_position_in_grid.x;
uint loc = elem_to_loc(elem, inp_shape, inp_strides, inp_ndim);
T tmp = inp[loc];
out[elem] = metal::exp(tmp);
out[elem] = metal::exp(tmp) * threads_per_simdgroup;
"""
source_contig = """
uint elem = thread_position_in_grid.x;
T tmp = inp[elem];
out[elem] = metal::exp(tmp);
out[elem] = metal::exp(tmp) * threads_per_simdgroup;
"""
# non contiguous
@@ -644,7 +644,7 @@ class TestFast(mlx_tests.MLXTestCase):
output_dtypes={"out": a.dtype},
stream=mx.gpu,
)
self.assertTrue(mx.allclose(mx.exp(a), outputs["out"]))
self.assertTrue(mx.allclose(mx.exp(a) * 32, outputs["out"]))
if __name__ == "__main__":