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

Fix round to round half-cases to even (#482)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-01-17 15:27:23 -08:00 committed by GitHub
parent 135fd796d2
commit 90c234b7ac
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 19 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/common/unary.h
View File

@ -56,11 +56,11 @@ struct SignOp {
struct RoundOp {
template <typename T>
T operator()(T x) {
return std::round(x);
return std::rint(x);
}
complex64_t operator()(complex64_t x) {
return {std::round(x.real()), std::round(x.imag())};
return {std::rint(x.real()), std::rint(x.imag())};
}
};

4
mlx/backend/metal/kernels/unary.metal
View File

@ -134,8 +134,8 @@ struct Negative {
};
struct Round {
template <typename T> T operator()(T x) { return metal::round(x); };
template <> complex64_t operator()(complex64_t x) { return {metal::round(x.real), metal::round(x.imag)}; };
template <typename T> T operator()(T x) { return metal::rint(x); };
template <> complex64_t operator()(complex64_t x) { return {metal::rint(x.real), metal::rint(x.imag)}; };
};
struct Sigmoid {

18
python/tests/test_ops.py
View File

@ -434,14 +434,14 @@ class TestOps(mlx_tests.MLXTestCase):
def test_round(self):
# float
x = mx.array(
[0.5, -0.5, 1.5, -1.5, -22.03, 19.98, -27, 9, 0.0, -np.inf, np.inf]
[0.5, -0.5, 1.5, -1.5, -21.03, 19.98, -27, 9, 0.0, -np.inf, np.inf]
)
expected = [1, -1, 2, -2, -22, 20, -27, 9, 0, -np.inf, np.inf]
expected = [0, -0, 2, -2, -21, 20, -27, 9, 0, -np.inf, np.inf]
self.assertListEqual(mx.round(x).tolist(), expected)
# complex
y = mx.round(mx.array([22.2 + 3.6j, 19.5 + 98.2j]))
self.assertListEqual(y.tolist(), [22 + 4j, 20 + 98j])
y = mx.round(mx.array([22.2 + 3.6j, 18.5 + 98.2j]))
self.assertListEqual(y.tolist(), [22 + 4j, 18 + 98j])
# decimals
y0 = mx.round(mx.array([15, 122], mx.int32), decimals=0)
@ -459,6 +459,16 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertTrue(mx.allclose(y1, mx.array([1.5, 1.5])))
self.assertTrue(mx.allclose(y2, mx.array([1.54, 1.47])))
# check round to nearest for different types
dtypes = [mx.bfloat16, mx.float16, mx.float32]
for dtype in dtypes:
x = mx.arange(10, dtype=dtype) - 4.5
x = mx.round(x)
self.assertEqual(
x.astype(mx.float32).tolist(),
[-4.0, -4.0, -2.0, -2.0, -0.0, 0.0, 2.0, 2.0, 4.0, 4.0],
)
def test_transpose_noargs(self):
x = mx.array([[0, 1, 1], [1, 0, 0]])

2
tests/ops_tests.cpp
View File

@ -957,7 +957,7 @@ TEST_CASE("test arithmetic unary ops") {
// Test round
{
array x({0.5, -0.5, 1.5, -1.5, 2.3, 2.6});
CHECK(array_equal(round(x), array({1, -1, 2, -2, 2, 3})).item<bool>());
CHECK(array_equal(round(x), array({0, -0, 2, -2, 2, 3})).item<bool>());
x = array({11, 222, 32});
CHECK(array_equal(round(x, -1), array({10, 220, 30})).item<bool>());