zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-01 12:49:44 +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

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") {