homogenize dipatching into pytato routines by characterizing into unary vs multi-ary

Author: kaushikcfd

arraycontext/impl/pytato/fake_numpy.py

@@ -50,20 +50,29 @@ class PytatoFakeNumpyNamespace(LoopyBasedFakeNumpyNamespace):
         :ref:`Pytato docs <pytato:memory-layout>` for more on this.
     """
 
-    _pt_funcs = frozenset({
+    _pt_unary_funcs = frozenset({
         "sin", "cos", "tan", "arcsin", "arccos", "arctan",
         "sinh", "cosh", "tanh", "exp", "log", "log10",
         "sqrt", "abs", "isnan", "real", "imag", "conj",
         })
 
+    _pt_multi_ary_funcs = frozenset({
+        "arctan2", "equal", "greater", "greater_equal", "less", "less_equal",
+        "not_equal", "minimum", "maximum", "where",
+    })
+
     def _get_fake_numpy_linalg_namespace(self):
         return PytatoFakeNumpyLinalgNamespace(self._array_context)
 
     def __getattr__(self, name):
-        if name in self._pt_funcs:
+        if name in self._pt_unary_funcs:
             from functools import partial
             return partial(rec_map_array_container, getattr(pt, name))
 
+        if name in self._pt_multi_ary_funcs:
+            from functools import partial
+            return partial(rec_multimap_array_container, getattr(pt, name))
+
         return super().__getattr__(name)
 
     # NOTE: the order of these follows the order in numpy docs
@@ -175,31 +184,10 @@ def rec_equal(x, y):
 
         return rec_equal(a, b)
 
-    def greater(self, x, y):
-        return rec_multimap_array_container(pt.greater, x, y)
-
-    def greater_equal(self, x, y):
-        return rec_multimap_array_container(pt.greater_equal, x, y)
-
-    def less(self, x, y):
-        return rec_multimap_array_container(pt.less, x, y)
-
-    def less_equal(self, x, y):
-        return rec_multimap_array_container(pt.less_equal, x, y)
-
-    def equal(self, x, y):
-        return rec_multimap_array_container(pt.equal, x, y)
-
-    def not_equal(self, x, y):
-        return rec_multimap_array_container(pt.not_equal, x, y)
-
     # }}}
 
     # {{{ mathematical functions
 
-    def arctan2(self, y, x):
-        return rec_multimap_array_container(pt.arctan2, y, x)
-
     def sum(self, a, axis=None, dtype=None):
         def _pt_sum(ary):
             if dtype not in [ary.dtype, None]:
@@ -209,18 +197,12 @@ def _pt_sum(ary):
 
         return rec_map_reduce_array_container(sum, _pt_sum, a)
 
-    def maximum(self, x, y):
-        return rec_multimap_array_container(pt.maximum, x, y)
-
     def amax(self, a, axis=None):
         return rec_map_reduce_array_container(
                 partial(reduce, pt.maximum), partial(pt.amax, axis=axis), a)
 
     max = amax
 
-    def minimum(self, x, y):
-        return rec_multimap_array_container(pt.minimum, x, y)
-
     def amin(self, a, axis=None):
         return rec_map_reduce_array_container(
                 partial(reduce, pt.minimum), partial(pt.amin, axis=axis), a)
@@ -231,10 +213,3 @@ def absolute(self, a):
         return self.abs(a)
 
     # }}}
-
-    # {{{ sorting, searching, and counting
-
-    def where(self, criterion, then, else_):
-        return rec_multimap_array_container(pt.where, criterion, then, else_)
-
-    # }}}