Improve, type, fix array_equal across all array contexts

Author: inducer

arraycontext/impl/jax/fake_numpy.py

@@ -27,12 +27,16 @@
 
 import jax.numpy as jnp
 
-from arraycontext.container import NotAnArrayContainerError, serialize_container
+from arraycontext.container import (
+    NotAnArrayContainerError,
+    serialize_container,
+)
 from arraycontext.container.traversal import (
     rec_map_array_container,
     rec_map_reduce_array_container,
     rec_multimap_array_container,
 )
+from arraycontext.context import Array, ArrayOrContainer
 from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace, BaseFakeNumpyNamespace
 
 
@@ -156,29 +160,35 @@ def any(self, a):
         return rec_map_reduce_array_container(
             partial(reduce, jnp.logical_or), jnp.any, a)
 
-    def array_equal(self, a, b):
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
         actx = self._array_context
 
         # NOTE: not all backends support `bool` properly, so use `int8` instead
-        true = actx.from_numpy(np.int8(True))
-        false = actx.from_numpy(np.int8(False))
+        true_ary = actx.from_numpy(np.int8(True))
+        false_ary = actx.from_numpy(np.int8(False))
 
         def rec_equal(x, y):
             if type(x) is not type(y):
-                return false
+                return false_ary
 
             try:
-                iterable = zip(serialize_container(x), serialize_container(y))
+                serialized_x = serialize_container(x)
+                serialized_y = serialize_container(y)
             except NotAnArrayContainerError:
                 if x.shape != y.shape:
-                    return false
+                    return false_ary
                 else:
                     return jnp.all(jnp.equal(x, y))
             else:
+                if len(serialized_x) != len(serialized_y):
+                    return false_ary
                 return reduce(
                         jnp.logical_and,
-                        [rec_equal(x_i, y_i) for (_, x_i), (_, y_i) in iterable],
-                        true)
+                        [(true_ary if kx_i == ky_i else false_ary)
+                            and rec_equal(x_i, y_i)
+                            for (kx_i, x_i), (ky_i, y_i)
+                            in zip(serialized_x, serialized_y)],
+                        true_ary)
 
         return rec_equal(a, b)
 

arraycontext/impl/numpy/fake_numpy.py

@@ -25,14 +25,14 @@
 
 import numpy as np
 
-from arraycontext.container import is_array_container
+from arraycontext.container import NotAnArrayContainerError, serialize_container
 from arraycontext.container.traversal import (
-    multimap_reduce_array_container,
     rec_map_array_container,
     rec_map_reduce_array_container,
     rec_multimap_array_container,
     rec_multimap_reduce_array_container,
 )
+from arraycontext.context import Array, ArrayOrContainer
 from arraycontext.fake_numpy import (
     BaseFakeNumpyLinalgNamespace,
     BaseFakeNumpyNamespace,
@@ -131,18 +131,34 @@ def all(self, a):
         return rec_map_reduce_array_container(partial(reduce, np.logical_and),
                                               lambda subary: np.all(subary), a)
 
-    def array_equal(self, a, b):
-        if type(a) != type(b):
-            return False
-        elif not is_array_container(a):
-            if a.shape != b.shape:
-                return False
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
+        def rec_equal(x: ArrayOrContainer, y: ArrayOrContainer) -> np.ndarray:
+            false_ary = np.array(False)
+            true_ary = np.array(True)
+            if type(x) is not type(y):
+                return false_ary
+
+            try:
+                serialized_x = serialize_container(x)
+                serialized_y = serialize_container(y)
+            except NotAnArrayContainerError:
+                assert isinstance(x, np.ndarray)
+                assert isinstance(y, np.ndarray)
+                return np.array(np.array_equal(x, y))
             else:
-                return np.all(np.equal(a, b))
-        else:
-            return multimap_reduce_array_container(partial(reduce,
-                                                           np.logical_and),
-                                                   self.array_equal, a, b)
+                if len(serialized_x) != len(serialized_y):
+                    return false_ary
+                return reduce(
+                        np.logical_and,
+                        [(true_ary if kx_i == ky_i else false_ary)
+                            and rec_equal(x_i, y_i)
+                            for (kx_i, x_i), (ky_i, y_i)
+                            in zip(serialized_x, serialized_y)],
+                        true_ary)
+
+        result = rec_equal(a, b)
+
+        return result
 
     def arange(self, *args, **kwargs):
         return np.arange(*args, **kwargs)

arraycontext/impl/pyopencl/fake_numpy.py

@@ -38,6 +38,7 @@
     rec_multimap_array_container,
     rec_multimap_reduce_array_container,
 )
+from arraycontext.context import Array, ArrayOrContainer
 from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace
 from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray
 from arraycontext.loopy import LoopyBasedFakeNumpyNamespace
@@ -215,30 +216,40 @@ def _any(ary):
             result = result.get()[()]
         return result
 
-    def array_equal(self, a, b):
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
         actx = self._array_context
         queue = actx.queue
 
         # NOTE: pyopencl doesn't like `bool` much, so use `int8` instead
-        true = actx.from_numpy(np.int8(True))
-        false = actx.from_numpy(np.int8(False))
+        true_ary = actx.from_numpy(np.int8(True))
+        false_ary = actx.from_numpy(np.int8(False))
 
-        def rec_equal(x, y):
+        def rec_equal(x: ArrayOrContainer, y: ArrayOrContainer) -> cl_array.Array:
             if type(x) is not type(y):
-                return false
+                return false_ary
 
             try:
-                iterable = zip(serialize_container(x), serialize_container(y))
+                serialized_x = serialize_container(x)
+                serialized_y = serialize_container(y)
             except NotAnArrayContainerError:
+                assert isinstance(x, cl_array.Array)
+                assert isinstance(y, cl_array.Array)
+
                 if x.shape != y.shape:
-                    return false
+                    return false_ary
                 else:
                     return (x == y).all()
             else:
+                if len(serialized_x) != len(serialized_y):
+                    return false_ary
+
                 return reduce(
                         partial(cl_array.minimum, queue=queue),
-                        [rec_equal(x_i, y_i)for (_, x_i), (_, y_i) in iterable],
-                        true)
+                        [(true_ary if kx_i == ky_i else false_ary)
+                            and rec_equal(x_i, y_i)
+                            for (kx_i, x_i), (ky_i, y_i)
+                            in zip(serialized_x, serialized_y)],
+                        true_ary)
 
         result = rec_equal(a, b)
         if not self._array_context._force_device_scalars:

arraycontext/impl/pytato/fake_numpy.py

@@ -22,7 +22,7 @@
 THE SOFTWARE.
 """
 from functools import partial, reduce
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 
@@ -34,6 +34,7 @@
     rec_map_reduce_array_container,
     rec_multimap_array_container,
 )
+from arraycontext.context import Array, ArrayOrContainer
 from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace
 from arraycontext.loopy import LoopyBasedFakeNumpyNamespace
 
@@ -171,31 +172,41 @@ def any(self, a):
                 partial(reduce, pt.logical_or),
                 lambda subary: pt.any(subary), a)
 
-    def array_equal(self, a, b):
+    def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
         actx = self._array_context
 
         # NOTE: not all backends support `bool` properly, so use `int8` instead
-        true = actx.from_numpy(np.int8(True))
-        false = actx.from_numpy(np.int8(False))
+        true_ary = actx.from_numpy(np.int8(True))
+        false_ary = actx.from_numpy(np.int8(False))
 
-        def rec_equal(x, y):
+        def rec_equal(x: ArrayOrContainer, y: ArrayOrContainer) -> pt.Array:
             if type(x) is not type(y):
-                return false
+                return false_ary
 
             try:
-                iterable = zip(serialize_container(x), serialize_container(y))
+                serialized_x = serialize_container(x)
+                serialized_y = serialize_container(y)
             except NotAnArrayContainerError:
+                assert isinstance(x, pt.Array)
+                assert isinstance(y, pt.Array)
+
                 if x.shape != y.shape:
-                    return false
+                    return false_ary
                 else:
-                    return pt.all(pt.equal(x, y))
+                    return pt.all(cast(pt.Array, pt.equal(x, y)))
             else:
+                if len(serialized_x) != len(serialized_y):
+                    return false_ary
+
                 return reduce(
                         pt.logical_and,
-                        [rec_equal(x_i, y_i) for (_, x_i), (_, y_i) in iterable],
-                        true)
+                        [(true_ary if kx_i == ky_i else false_ary)
+                            and rec_equal(x_i, y_i)
+                            for (kx_i, x_i), (ky_i, y_i)
+                            in zip(serialized_x, serialized_y)],
+                        true_ary)
 
-        return rec_equal(a, b)
+        return cast(Array, rec_equal(a, b))
 
     # }}}