fix numpy test failing

ASEM000 · Mar 27, 2024 · 3f4b5dd · 3f4b5dd
1 parent 173aaf7
commit 3f4b5dd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/sepes/_src/backend/arraylib/__init__.py b/sepes/_src/backend/arraylib/__init__.py
@@ -22,8 +22,9 @@
 class NoImplError(NamedTuple):
     op: Callable
 
-    def __call__(self, *_, **__):
-        raise NotImplementedError(f"No implementation for {self.op}")
+    def __call__(self, *args, **kwargs):
+        raise NotImplementedError(f"No implementation for {self.op}"
+                                  f" with {args=} {kwargs=}")
 
 
 tobytes = ft.singledispatch(NoImplError("tobytes"))

diff --git a/sepes/_src/tree_util.py b/sepes/_src/tree_util.py
@@ -82,14 +82,17 @@ def is_array_like(node) -> bool:
     return hasattr(node, "shape") and hasattr(node, "dtype")
 
 
-def _is_leaf_rhs_equal(leaf, rhs) -> bool:
+def _is_leaf_rhs_equal(leaf, rhs):
     if is_array_like(leaf):
         if is_array_like(rhs):
             if leaf.shape != rhs.shape:
                 return False
             if leaf.dtype != rhs.dtype:
                 return False
-            verdict = arraylib.all(leaf == rhs)
+            try:
+                verdict = arraylib.all(leaf == rhs)
+            except NotImplementedError:
+                verdict = leaf == rhs
             try:
                 return bool(verdict)
             except Exception: