From c677cc9e744ae5df000527433d76f2172cdd434f Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 28 Dec 2021 15:24:22 +0100 Subject: [PATCH 1/6] Unify pytato reduction handling, support 'initial' arg --- arraycontext/impl/pytato/fake_numpy.py | 54 ++++++++++++++++++++------ 1 file changed, 43 insertions(+), 11 deletions(-) diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index 905fd0f8..17173a94 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -42,6 +42,10 @@ class PytatoFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace): pass +class _NoValue: + pass + + class PytatoFakeNumpyNamespace(BaseFakeNumpyNamespace): """ A :mod:`numpy` mimic for :class:`PytatoPyOpenCLArrayContext`. @@ -91,22 +95,50 @@ def minimum(self, x, y): def where(self, criterion, then, else_): return rec_multimap_array_container(pt.where, criterion, then, else_) - def sum(self, a, axis=None, dtype=None): - def _pt_sum(ary): + @staticmethod + def _reduce(container_binop, array_reduce, + ary, *, + axis, dtype, initial): + def container_reduce(ctr): + if initial is _NoValue: + try: + return reduce(container_binop, ctr) + except TypeError as exc: + assert "empty sequence" in str(exc) + raise ValueError("zero-size reduction operation " + "without supplied 'initial' value") + else: + return reduce(container_binop, ctr, initial) + + def actual_array_reduce(ary): if dtype not in [ary.dtype, None]: raise NotImplementedError - return pt.sum(ary, axis=axis) - - return rec_map_reduce_array_container(sum, _pt_sum, a) - - def min(self, a, axis=None): - return rec_map_reduce_array_container( - partial(reduce, pt.minimum), partial(pt.amin, axis=axis), a) + if initial is _NoValue: + return array_reduce(ary, axis=axis) + else: + return array_reduce(ary, axis=axis, initial=initial) - def max(self, a, axis=None): return rec_map_reduce_array_container( - partial(reduce, pt.maximum), partial(pt.amax, axis=axis), a) + container_reduce, + actual_array_reduce, + ary) + + # * appears where positional signature starts diverging from numpy + def sum(self, a, axis=None, dtype=None, *, initial=0): + import operator + return self._reduce(operator.add, pt.sum, a, + axis=axis, dtype=dtype, initial=initial) + + # * appears where positional signature starts diverging from numpy + def min(self, a, axis=None, *, initial=_NoValue): + return self._reduce(pt.minimum, pt.amin, a, + axis=axis, dtype=None, initial=initial) + + # * appears where positional signature starts diverging from numpy + def max(self, a, axis=None, *, initial=_NoValue): + return self._reduce(pt.maximum, pt.amax, a, + axis=axis, dtype=None, initial=initial) def stack(self, arrays, axis=0): return rec_multimap_array_container( From 4c06f61474064f8dce75682c039e3dd0a811a920 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andreas=20Kl=C3=B6ckner?= Date: Thu, 6 Jan 2022 01:58:58 -0600 Subject: [PATCH 2/6] Pass on unrelated TypeErrors in pytato reduction handling Co-authored-by: Matt Smith --- arraycontext/impl/pytato/fake_numpy.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index 17173a94..51e0bf19 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -104,9 +104,11 @@ def container_reduce(ctr): try: return reduce(container_binop, ctr) except TypeError as exc: - assert "empty sequence" in str(exc) - raise ValueError("zero-size reduction operation " - "without supplied 'initial' value") + if "empty sequence" in str(exc): + raise ValueError("zero-size reduction operation " + "without supplied 'initial' value") + else: + raise else: return reduce(container_binop, ctr, initial) From e23a680d37974f61627b5bb4dd90ad007bad3dfc Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 6 Jan 2022 10:32:23 -0600 Subject: [PATCH 3/6] flake8 --- arraycontext/impl/pytato/fake_numpy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index 51e0bf19..44c5f9bc 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -108,7 +108,7 @@ def container_reduce(ctr): raise ValueError("zero-size reduction operation " "without supplied 'initial' value") else: - raise + raise else: return reduce(container_binop, ctr, initial) From 1f15ae9332d7f24051064c213ce452d8cab8a5eb Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 6 Jan 2022 12:10:23 -0600 Subject: [PATCH 4/6] explain _NoValue --- arraycontext/impl/pytato/fake_numpy.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index 44c5f9bc..37456bc6 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -133,11 +133,15 @@ def sum(self, a, axis=None, dtype=None, *, initial=0): axis=axis, dtype=dtype, initial=initial) # * appears where positional signature starts diverging from numpy + # Use _NoValue to indicate a lack of neutral element so that the caller can + # use None as a neutral element def min(self, a, axis=None, *, initial=_NoValue): return self._reduce(pt.minimum, pt.amin, a, axis=axis, dtype=None, initial=initial) # * appears where positional signature starts diverging from numpy + # Use _NoValue to indicate a lack of neutral element so that the caller can + # use None as a neutral element def max(self, a, axis=None, *, initial=_NoValue): return self._reduce(pt.maximum, pt.amax, a, axis=axis, dtype=None, initial=initial) From d4271f22b30895086d001ed09030da36efc58961 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 6 Jan 2022 12:10:56 -0600 Subject: [PATCH 5/6] cosmetic changes --- arraycontext/impl/pytato/fake_numpy.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index 37456bc6..b4b21ba7 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -96,13 +96,13 @@ def where(self, criterion, then, else_): return rec_multimap_array_container(pt.where, criterion, then, else_) @staticmethod - def _reduce(container_binop, array_reduce, + def _map_reduce(reduce_func, map_func, ary, *, axis, dtype, initial): - def container_reduce(ctr): + def _reduce_func(partial_results): if initial is _NoValue: try: - return reduce(container_binop, ctr) + return reduce(reduce_func, partial_results) except TypeError as exc: if "empty sequence" in str(exc): raise ValueError("zero-size reduction operation " @@ -110,40 +110,40 @@ def container_reduce(ctr): else: raise else: - return reduce(container_binop, ctr, initial) + return reduce(reduce_func, partial_results, initial) - def actual_array_reduce(ary): - if dtype not in [ary.dtype, None]: + def _map_func(ary): + if dtype not in [None, ary.dtype]: raise NotImplementedError if initial is _NoValue: - return array_reduce(ary, axis=axis) + return map_func(ary, axis=axis) else: - return array_reduce(ary, axis=axis, initial=initial) + return map_func(ary, axis=axis, initial=initial) return rec_map_reduce_array_container( - container_reduce, - actual_array_reduce, + _reduce_func, + _map_func, ary) # * appears where positional signature starts diverging from numpy def sum(self, a, axis=None, dtype=None, *, initial=0): import operator - return self._reduce(operator.add, pt.sum, a, + return self._map_reduce(operator.add, pt.sum, a, axis=axis, dtype=dtype, initial=initial) # * appears where positional signature starts diverging from numpy # Use _NoValue to indicate a lack of neutral element so that the caller can # use None as a neutral element def min(self, a, axis=None, *, initial=_NoValue): - return self._reduce(pt.minimum, pt.amin, a, + return self._map_reduce(pt.minimum, pt.amin, a, axis=axis, dtype=None, initial=initial) # * appears where positional signature starts diverging from numpy # Use _NoValue to indicate a lack of neutral element so that the caller can # use None as a neutral element def max(self, a, axis=None, *, initial=_NoValue): - return self._reduce(pt.maximum, pt.amax, a, + return self._map_reduce(pt.maximum, pt.amax, a, axis=axis, dtype=None, initial=initial) def stack(self, arrays, axis=0): From 994918ab4bd07e040d3cb2384cd23e44b0dfcf17 Mon Sep 17 00:00:00 2001 From: Matthew Smith Date: Thu, 6 Jan 2022 12:12:32 -0600 Subject: [PATCH 6/6] add empty container support to pyopencl's np.[sum,min,max] --- arraycontext/impl/pyopencl/fake_numpy.py | 95 +++++++++++++++--------- 1 file changed, 61 insertions(+), 34 deletions(-) diff --git a/arraycontext/impl/pyopencl/fake_numpy.py b/arraycontext/impl/pyopencl/fake_numpy.py index 4e9d48be..7d1de932 100644 --- a/arraycontext/impl/pyopencl/fake_numpy.py +++ b/arraycontext/impl/pyopencl/fake_numpy.py @@ -50,6 +50,10 @@ # {{{ fake numpy +class _NoValue: + pass + + class PyOpenCLFakeNumpyNamespace(BaseFakeNumpyNamespace): def _get_fake_numpy_linalg_namespace(self): return _PyOpenCLFakeNumpyLinalgNamespace(self._array_context) @@ -110,58 +114,81 @@ def where_inner(inner_crit, inner_then, inner_else): return rec_multimap_array_container(where_inner, criterion, then, else_) - def sum(self, a, axis=None, dtype=None): - + @staticmethod + def _map_reduce(reduce_func, map_func, + ary, *, + axis, dtype, initial): if isinstance(axis, int): axis = axis, - def _rec_sum(ary): + def _reduce_func(partial_results): + if initial is _NoValue: + try: + return reduce(reduce_func, partial_results) + except TypeError as exc: + if "empty sequence" in str(exc): + raise ValueError("zero-size reduction operation " + "without supplied 'initial' value") + else: + raise + else: + return reduce(reduce_func, partial_results, initial) + + def _map_func(ary): + if dtype not in [None, ary.dtype]: + raise NotImplementedError if axis not in [None, tuple(range(ary.ndim))]: - raise NotImplementedError(f"Sum over '{axis}' axes not supported.") + raise NotImplementedError( + f"Mapping over '{axis}' axes not supported.") + if initial is _NoValue: + return map_func(ary) + else: + return map_func(ary, initial=initial) - return cl_array.sum(ary, dtype=dtype, queue=self._array_context.queue) + return rec_map_reduce_array_container( + _reduce_func, + _map_func, + ary) - result = rec_map_reduce_array_container(sum, _rec_sum, a) + # * appears where positional signature starts diverging from numpy + def sum(self, a, axis=None, dtype=None, *, initial=0): + queue = self._array_context.queue + import operator + result = self._map_reduce( + operator.add, + partial(cl_array.sum, queue=queue), + a, + axis=axis, dtype=dtype, initial=initial) if not self._array_context._force_device_scalars: result = result.get()[()] return result - def min(self, a, axis=None): + # * appears where positional signature starts diverging from numpy + # Use _NoValue to indicate a lack of neutral element so that the caller can + # use None as a neutral element + def min(self, a, axis=None, *, initial=_NoValue): queue = self._array_context.queue - - if isinstance(axis, int): - axis = axis, - - def _rec_min(ary): - if axis not in [None, tuple(range(ary.ndim))]: - raise NotImplementedError(f"Min. over '{axis}' axes not supported.") - return cl_array.min(ary, queue=queue) - - result = rec_map_reduce_array_container( - partial(reduce, partial(cl_array.minimum, queue=queue)), - _rec_min, - a) + result = self._map_reduce( + partial(cl_array.minimum, queue=queue), + partial(cl_array.min, queue=queue), + a, + axis=axis, dtype=None, initial=initial) if not self._array_context._force_device_scalars: result = result.get()[()] return result - def max(self, a, axis=None): + # * appears where positional signature starts diverging from numpy + # Use _NoValue to indicate a lack of neutral element so that the caller can + # use None as a neutral element + def max(self, a, axis=None, *, initial=_NoValue): queue = self._array_context.queue - - if isinstance(axis, int): - axis = axis, - - def _rec_max(ary): - if axis not in [None, tuple(range(ary.ndim))]: - raise NotImplementedError(f"Max. over '{axis}' axes not supported.") - return cl_array.max(ary, queue=queue) - - result = rec_map_reduce_array_container( - partial(reduce, partial(cl_array.maximum, queue=queue)), - _rec_max, - a) + result = self._map_reduce( + partial(cl_array.maximum, queue=queue), + partial(cl_array.max, queue=queue), + a, + axis=axis, dtype=None, initial=initial) if not self._array_context._force_device_scalars: result = result.get()[()]