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

recompile when shapeless is different (#1776)

Browse Source
This commit is contained in:
Awni Hannun 2025-01-20 21:07:10 -08:00 committed by GitHub
parent a8666a757a
commit 90532b1f37
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlx/compile.cpp
View File

@ -197,8 +197,10 @@ std::uintptr_t get_function_address(const std::function<T(U...)>& fun) {
class CompilerCache {
public:
struct CacheEntry {
CacheEntry(Stream stream) : stream(stream) {};
CacheEntry(Stream stream, bool shapeless)
: stream(stream), shapeless(shapeless) {};
Stream stream;
bool shapeless;
std::vector<array> inputs;
std::vector<array> outputs;
std::vector<array> tape;
@ -245,6 +247,9 @@ class CompilerCache {
if (entry.stream != stream) {
continue;
}
if (entry.shapeless != shapeless) {
continue;
}
// Check the inputs match and return if so
if (has_same_shape_and_dtype(inputs, entry.inputs) &&
@ -253,7 +258,7 @@ class CompilerCache {
}
}
// Otherwise append a new cache entry
entries.push_back(CacheEntry{stream});
entries.push_back(CacheEntry{stream, shapeless});
return entries.back();
}

4
python/tests/test_compile.py
View File

@ -675,6 +675,10 @@ class TestCompile(mlx_tests.MLXTestCase):
def mean(x):
return mx.mean(x, keepdims=True)
cfun = mx.compile(mean)
out = cfun(mx.ones((5, 5)))
self.assertTrue(mx.allclose(out, mx.array(1.0)))
cmean = mx.compile(mean, shapeless=True)
x = mx.ones(2)