L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Option to JIT steel gemm / conv (#1139)

Browse Source
This commit is contained in:
Awni Hannun
2024-05-23 18:07:34 -07:00
committed by GitHub
parent eab2685c67
commit 7e26fd8032
31 changed files with 2504 additions and 1540 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/src/ops.cpp
+3 -3
View File
@@ -3968,9 +3968,9 @@ void init_ops(nb::module_& m) {
a (array): Input array or scalar.
b (array): Input array or scalar.
block_size (int): Size of blocks to be masked. Must be ``32`` or ``64`` (default: ``64``)
mask_out (array, optional): Boolean mask for output (default: ``None``)
mask_lhs (array, optional): Boolean mask for a (default: ``None``)
mask_rhs (array, optional): Boolean mask for b (default: ``None``)
mask_out (array, optional): Mask for output (default: ``None``)
mask_lhs (array, optional): Mask for a (default: ``None``)
mask_rhs (array, optional): Mask for b (default: ``None``)
)pbdoc");
m.def(
python/tests/test_blas.py
+1 -3
View File
@@ -556,8 +556,7 @@ class TestBlas(mlx_tests.MLXTestCase):
self.assertListEqual(list(d_npy.shape), list(d_mlx.shape))
self.assertTrue(np.allclose(d_mlx, d_npy, atol=1e-5))
# # Batched matmul with simple broadcast
# Batched matmul with simple broadcast
a_npy = np.random.normal(0.0, 1.0 / 128, (32, 128, 16)).astype(np.float32)
b_npy = np.random.normal(0.0, 1.0 / 128, (16, 16)).astype(np.float32)
@@ -573,7 +572,6 @@ class TestBlas(mlx_tests.MLXTestCase):
self.assertListEqual(list(d_npy.shape), list(d_mlx.shape))
self.assertTrue(np.allclose(d_mlx, d_npy, atol=1e-5))
# Matmul with vector
a_npy = np.random.normal(0.0, 1.0 / 128, (16,)).astype(np.float32)
b_npy = np.random.normal(0.0, 1.0 / 128, (32, 16, 128)).astype(np.float32)