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

fix grad, more tests

Browse Source
This commit is contained in:
Awni Hannun
2024-01-14 19:42:36 -08:00
parent 9d70412a91
commit 9739c72781
2 changed files with 47 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
tests/compile_tests.cpp
View File

@@ -8,13 +8,42 @@ using namespace mlx::core;
std::vector<array> simple_fun(const std::vector<array>& inputs) {
return std::vector<array>{inputs[0] + inputs[1]};
};
}
TEST_CASE("test simple compile") {
auto compfn = compile(simple_fun);
auto out = compfn({array(1.0), array(2.0)})[0];
auto out = compfn({array(1.0f), array(2.0f)})[0];
CHECK_EQ(out.item<float>(), 3.0f);
out = compfn({array(1.0), array(2.0)})[0];
out = compfn({array(1.0f), array(2.0f)})[0];
CHECK_EQ(out.item<float>(), 3.0f);
// Change the shapes
out = compfn({array({1.0f, 2.0f}), array(2.0f)})[0];
CHECK(array_equal(out, array({3.0f, 4.0f})).item<bool>());
out = compfn({array(2.0f), array({1.0f, 2.0f})})[0];
CHECK(array_equal(out, array({3.0f, 4.0f})).item<bool>());
// Change the types
out = compfn({array(2, int32), array({1.0f, 2.0f})})[0];
CHECK(array_equal(out, array({3.0f, 4.0f})).item<bool>());
out = compfn({array(2.0f), array({1, 2}, int32)})[0];
CHECK(array_equal(out, array({3.0f, 4.0f})).item<bool>());
}
std::vector<array> fun1(const std::vector<array>& inputs) {
auto loss = [](std::vector<array> ins) { return exp(ins[0] + ins[1]); };
return grad(loss)(inputs);
}
TEST_CASE("test compile with grad") {
auto x = array(1.0f);
auto y = array(1.0f);
auto grads_expected = fun1({x, y});
auto grads_compile = compile(fun1)({x, y});
CHECK_EQ(grads_compile[0].item<float>(), grads_expected[0].item<float>());
}
TEST_CASE("test nested compile") {}