From 32eae198c227c3691f28a7828b4208f69970fb38 Mon Sep 17 00:00:00 2001
From: ohrechykha <alexeee6@gmail.com>
Date: Thu, 8 Aug 2024 13:01:27 +0300
Subject: [PATCH] implementing a wrapper to fix test errors

---
 src/ragged/_spec_elementwise_functions.py | 14 ++++++++------
 tests/test_spec_elementwise_functions.py  | 19 +++++++++++++++----
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/ragged/_spec_elementwise_functions.py b/src/ragged/_spec_elementwise_functions.py
index 0905e9a..654f98e 100644
--- a/src/ragged/_spec_elementwise_functions.py
+++ b/src/ragged/_spec_elementwise_functions.py
@@ -413,16 +413,18 @@ def ceil(x: array, /) -> array:
 
     https://data-apis.org/array-api/latest/API_specification/generated/array_api.ceil.html
     """
-    def _wrapper(dtype):
-        if dtype in [np.int8, np.uint8, np.bool_, np.bool]:
+
+    def _wrapper(t: np.dtype, /) -> np.dtype:
+        if t in [np.int8, np.uint8, np.bool_, bool]:
             return np.float16
-        elif dtype in [np.int16, np.uint16]:
+        elif t in [np.int16, np.uint16]:
             return np.float32
-        elif dtype in [np.int32, np.uint32, np.int64, np.uint64]:
+        elif t in [np.int32, np.uint32, np.int64, np.uint64]:
             return np.float64
         else:
-            return dtype
-    return _box(type(x), np.ceil(*_unbox(x)), _wrapper(x.dtype))
+            return t
+
+    return _box(type(x), np.ceil(*_unbox(x)), dtype=_wrapper(x.dtype))
 
 
 def conj(x: array, /) -> array:
diff --git a/tests/test_spec_elementwise_functions.py b/tests/test_spec_elementwise_functions.py
index 7afebd2..d1e3477 100644
--- a/tests/test_spec_elementwise_functions.py
+++ b/tests/test_spec_elementwise_functions.py
@@ -49,6 +49,17 @@ def first(x: ragged.array) -> Any:
     return xp.asarray(out.item(), dtype=x.dtype)
 
 
+def _wrapper(t: np.dtype, /) -> np.dtype:
+    if t in [np.int8, np.uint8, np.bool_, bool]:
+        return np.float16
+    elif t in [np.int16, np.uint16]:
+        return np.float32
+    elif t in [np.int32, np.uint32, np.int64, np.uint64]:
+        return np.float64
+    else:
+        return t
+
+
 def test_existence():
     assert ragged.abs is not None
     assert ragged.acos is not None
@@ -394,8 +405,8 @@ def test_ceil_int(device, x_int):
     result = ragged.ceil(x_int.to_device(device))
     assert type(result) is type(x_int)
     assert result.shape == x_int.shape
-    assert xp.ceil(first(x_int)) == first(result)
-    assert xp.ceil(first(x_int)).dtype == result.dtype
+    assert xp.ceil(first(x_int)) == first(result).astype(_wrapper(first(result).dtype))
+    assert xp.ceil(first(x_int)).dtype == _wrapper(result.dtype)
 
 
 @pytest.mark.skipif(
@@ -507,8 +518,8 @@ def test_floor_int(device, x_int):
     result = ragged.floor(x_int.to_device(device))
     assert type(result) is type(x_int)
     assert result.shape == x_int.shape
-    assert xp.floor(first(x_int)) == first(result)
-    assert xp.floor(first(x_int)).dtype == result.dtype
+    assert xp.floor(first(x_int)) == first(result).astype(_wrapper(first(result).dtype))
+    assert xp.floor(first(x_int)).dtype == _wrapper(result.dtype)
 
 
 @pytest.mark.parametrize("device", devices)