zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-21 10:18:10 +08:00
Code Issues 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

6
python/src/ops.cpp
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(