zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-03 22:34:43 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

compile changes if stream changes (#1644)

Browse Source
This commit is contained in:
Awni Hannun
2024-12-03 14:37:44 -08:00
committed by GitHub
parent 9d40e521d7
commit e047fd977d
4 changed files with 24 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/compile_tests.cpp
View File

@@ -719,3 +719,14 @@ TEST_CASE("test compile strides") {
CHECK_EQ(out.strides().size(), 3);
}
}
TEST_CASE("test compile change streams") {
auto cfun = compile(simple_fun);
auto out = cfun({array(1.0f), array(2.0f)})[0];
CHECK_EQ(out.primitive().stream(), default_stream(default_device()));
auto s = new_stream(default_device());
StreamContext sctx(s);
out = cfun({array(1.0f), array(2.0f)})[0];
CHECK_EQ(out.primitive().stream(), s);
}