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 1 Packages Projects Releases Wiki Activity

NumberOfElements for shapeless compile and vmap fixes (#802)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-03-13 10:34:14 -07:00
committed by GitHub
parent 29d0c10ee5
commit 76c919b4ec
13 changed files with 289 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mlx/compile.cpp
View File

@@ -74,7 +74,7 @@ bool allows_shapeless(const Primitive& p) {
is_noop(p) || is_reduction(p) || typeid(p) == typeid(Softmax) ||
typeid(p) == typeid(Sort) || typeid(p) == typeid(ArgSort) ||
typeid(p) == typeid(ArgPartition) || typeid(p) == typeid(Partition) ||
typeid(p) == typeid(Select);
typeid(p) == typeid(Select) || typeid(p) == typeid(NumberOfElements);
}
Compiled::Compiled(