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

Implement the 'where' primitive for conditional selection (#664)

Browse Source
This commit is contained in:
Rifur13
2024-02-22 18:10:48 -05:00
committed by GitHub
parent ad4a45e615
commit 126c9869c8
23 changed files with 991 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
tests/autograd_tests.cpp
View File

@@ -1075,6 +1075,37 @@ TEST_CASE("test jvp from vjp") {
CHECK(compute_derivs(subtract));
CHECK(compute_derivs(power));
}
// Conditional selection element-wise op
{
auto condition = random::randint(0, 2, {5, 10});
auto x = random::uniform({5, 10});
auto y = random::uniform({5, 10});
eval(condition, x, y);
auto compute_derivs = [&condition, &x, &y](auto fn) {
auto fn_wrap = [&fn](std::vector<array> inputs) {
return std::vector<array>{
fn(inputs[0], inputs[1], inputs[2], default_device())};
};
// Compute vjp and add results
auto vjps = vjp(fn_wrap, {condition, x, y}, {ones(x.shape())}).second;
auto vjp_out = add(add(vjps[0], vjps[1]), vjps[2]);
// Compute jvp
array jvp_out =
jvp(fn_wrap,
{condition, x, y},
{ones(condition.shape()), ones(y.shape()), ones(x.shape())})
.second[0];
array result = array_equal(vjp_out, jvp_out);
return result.item<bool>();
};
CHECK(compute_derivs(where));
}
}
TEST_CASE("test complex gradients") {

45
tests/ops_tests.cpp
View File

@@ -2193,6 +2193,8 @@ TEST_CASE("test power") {
}
TEST_CASE("test where") {
const float inf = std::numeric_limits<float>::infinity();
array condition(true);
array x(1.0f);
array y(0.0f);
@@ -2224,6 +2226,49 @@ TEST_CASE("test where") {
out = where(condition, x, y);
expected = array({1, 2, 2, 1}, {2, 2});
CHECK(array_equal(where(condition, x, y), expected).item<bool>());
condition = array(true);
x = array({1, 2, 3});
y = array({3, 6, 13});
CHECK(array_equal(where(condition, x, y), array({1, 2, 3})).item<bool>());
condition = array(false);
x = array({1, 2, 3});
y = array({3, 6, 13});
CHECK(array_equal(where(condition, x, y), array({3, 6, 13})).item<bool>());
condition = array({1, 1, 0});
x = array({1, 2, 3});
y = array({11, 12, 13});
CHECK(array_equal(where(condition, x, y), array({1, 2, 13})).item<bool>());
condition = array({true, false}, {2, 1, 1});
x = array({1, 2, 3, 4}, {2, 1, 2});
y = array({11, 22, 33, 44}, {2, 2, 1});
expected = array({1, 2, 1, 2, 33, 33, 44, 44}, {2, 2, 2});
CHECK(array_equal(where(condition, x, y), expected).item<bool>());
condition = array({true, false, false});
x = array({inf, 2.0, 3.0});
y = array({10.0, 20.0, -inf});
CHECK(array_equal(where(condition, x, y), array({inf, 20.0, -inf}))
.item<bool>());
// 4-dim optimized case.
condition = array({false});
x = array({1, 2}, {2, 1, 1, 1});
y = array({3, 4}, {1, 1, 2, 1});
CHECK(array_equal(where(condition, x, y), array({3, 4, 3, 4}, {2, 1, 2, 1}))
.item<bool>());
// 5-dim optimized case.
condition = array({true, false}, {2, 1, 1, 1, 1});
x = array({1, 2, 3, 4}, {2, 1, 1, 1, 2});
y = array({11, 22}, {1, 1, 2, 1, 1});
CHECK(array_equal(
where(condition, x, y),
array({1, 2, 1, 2, 11, 11, 22, 22}, {2, 1, 2, 1, 2}))
.item<bool>());
}
TEST_CASE("test stack") {

64
tests/vmap_tests.cpp
View File

@@ -138,6 +138,70 @@ TEST_CASE("test simple vmap") {
CHECK(array_equal(out, x + y).item<bool>());
}
// vmap where (ternary op)
{
auto fun = [](std::vector<array> inputs) {
auto out = where(inputs[0], inputs[1], inputs[2]);
return std::vector<array>{out};
};
auto vfun = vmap(fun);
array cond({true, false}, {2, 1});
array x({1.0, 2.0}, {2, 1});
array y({2.0, 4.0}, {2, 1});
auto out = vfun({cond, x, y})[0];
CHECK(array_equal(out, array({1.0, 4.0}, {2, 1})).item<bool>());
cond = array({true, true, false}, {1, 3});
x = ones({2, 1, 3});
y = zeros({3, 2});
vfun = vmap(fun, {1, 2, 0});
out = vfun({cond, x, y})[0];
CHECK(
array_equal(out, array({1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0}, {3, 2, 2}))
.item<bool>());
vfun = vmap(fun, {1, 2, 0}, {1});
out = vfun({cond, x, y})[0];
CHECK(
array_equal(out, array({1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0}, {2, 3, 2}))
.item<bool>());
cond = array({true, false});
x = array(2.);
y = ones({3, 2});
vfun = vmap(fun, {-1, -1, 0});
out = vfun({cond, x, y})[0];
CHECK(array_equal(out, array({2, 1, 2, 1, 2, 1}, {3, 2})).item<bool>());
cond = array({true, false});
x = ones({3, 2});
y = array(2.);
vfun = vmap(fun, {-1, 0, -1});
out = vfun({cond, x, y})[0];
CHECK(array_equal(out, array({1, 2, 1, 2, 1, 2}, {3, 2})).item<bool>());
CHECK_THROWS_AS(vmap(fun, {-1, -1, -1}, {0}), std::invalid_argument);
CHECK_THROWS_AS(vmap(fun, {-1, 0, -1}, {-1}), std::invalid_argument);
CHECK_THROWS_AS(vmap(fun, {-1, -1, 0}, {-1}), std::invalid_argument);
CHECK_THROWS_AS(vmap(fun, {0, -1, -1}, {-1}), std::invalid_argument);
cond = array({true, false});
x = array(1.);
y = array(2.);
vfun = vmap(fun, {-1, -1, -1}, {-1});
out = vfun({cond, x, y})[0];
CHECK(array_equal(out, array({1.0, 2.0})).item<bool>());
cond = array({1, 1, 1, 0, 0, 0}, {3, 2, 1});
x = ones({3, 2, 1});
y = full({3, 2, 1}, 2);
vfun = vmap(vmap(fun));
out = vfun({cond, x, y})[0];
CHECK(array_equal(out, array({1, 1, 1, 2, 2, 2}, {3, 2, 1})).item<bool>());
}
// vmap with capturing closure
{
auto x = add(add(ones({2}), zeros({2})), zeros({2}));