Merge branch 'main' into pytato-array-context-transforms

Author: thomasgibson

arraycontext/container/__init__.py

@@ -22,7 +22,7 @@
 
 Serialization/deserialization
 -----------------------------
-.. autofunction:: is_array_container
+.. autofunction:: is_array_container_type
 .. autofunction:: serialize_container
 .. autofunction:: deserialize_container
 
@@ -95,8 +95,6 @@ class ArrayContainer:
 
     This allows enumeration of the component arrays in a container and the
     construction of modified containers from an iterable of those component arrays.
-    :func:`is_array_container` will return *True* for types that have
-    a container serialization function registered.
 
     Packages may register their own types as array containers. They must not
     register other types (e.g. :class:`list`) as array containers.
@@ -156,6 +154,13 @@ def is_array_container_type(cls: type) -> bool:
     """
     :returns: *True* if the type *cls* has a registered implementation of
         :func:`serialize_container`, or if it is an :class:`ArrayContainer`.
+
+    .. warning::
+
+        Not all instances of a type that this function labels an array container
+        must automatically be array containers. For example, while this
+        function will say that :class:`numpy.ndarray` is an array container
+        type, only object arrays *actually are* array containers.
     """
     return (
             cls is ArrayContainer
@@ -168,6 +173,13 @@ def is_array_container(ary: Any) -> bool:
     :returns: *True* if the instance *ary* has a registered implementation of
         :func:`serialize_container`.
     """
+
+    from warnings import warn
+    warn("is_array_container is deprecated and will be removed in 2022. "
+            "If you must know precisely whether something is an array container, "
+            "try serializing it and catch TypeError. For a cheaper option, see "
+            "is_array_container_type.",
+            DeprecationWarning, stacklevel=2)
     return (serialize_container.dispatch(ary.__class__)
             is not serialize_container.__wrapped__)       # type:ignore[attr-defined]
 
@@ -190,7 +202,7 @@ def get_container_context(ary: ArrayContainer) -> Optional[ArrayContext]:
 @serialize_container.register(np.ndarray)
 def _serialize_ndarray_container(ary: np.ndarray) -> Iterable[Tuple[Any, Any]]:
     if ary.dtype.char != "O":
-        raise ValueError(
+        raise TypeError(
                 f"cannot seriealize '{type(ary).__name__}' with dtype '{ary.dtype}'")
 
     # special-cased for speed
@@ -232,7 +244,7 @@ def get_container_context_recursively(ary: Any) -> Optional[ArrayContext]:
     any level, an assertion error is raised.
     """
     actx = None
-    if not is_array_container(ary):
+    if not is_array_container_type(ary.__class__):
         return actx
 
     # try getting the array context directly

arraycontext/container/dataclass.py

@@ -44,7 +44,7 @@ def dataclass_array_container(cls: type) -> type:
 
     Attributes that are not array containers are allowed. In order to decide
     whether an attribute is an array container, the declared attribute type
-    is checked by the criteria from :func:`is_array_container`.
+    is checked by the criteria from :func:`is_array_container_type`.
     """
     from dataclasses import is_dataclass
     assert is_dataclass(cls)

arraycontext/container/traversal.py

@@ -60,7 +60,7 @@
 
 from arraycontext.context import ArrayContext
 from arraycontext.container import (
-        ContainerT, ArrayOrContainerT, is_array_container,
+        ContainerT, ArrayOrContainerT, is_array_container_type,
         serialize_container, deserialize_container)
 
 
@@ -81,7 +81,7 @@ def _map_array_container_impl(
     def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
         if type(_ary) is leaf_cls:  # type(ary) is never None
             return f(_ary)
-        elif is_array_container(_ary):
+        elif is_array_container_type(_ary.__class__):
             return deserialize_container(_ary, [
                 (key, frec(subary)) for key, subary in serialize_container(_ary)
                 ])
@@ -108,7 +108,7 @@ def _multimap_array_container_impl(
     def rec(*_args: Any) -> Any:
         template_ary = _args[container_indices[0]]
         if (type(template_ary) is leaf_cls
-                or not is_array_container(template_ary)):
+                or not is_array_container_type(template_ary.__class__)):
             return f(*_args)
 
         assert all(
@@ -136,7 +136,7 @@ def rec(*_args: Any) -> Any:
 
     container_indices: List[int] = [
             i for i, arg in enumerate(args)
-            if is_array_container(arg) and type(arg) is not leaf_cls]
+            if is_array_container_type(arg.__class__) and type(arg) is not leaf_cls]
 
     if not container_indices:
         return f(*args)
@@ -448,7 +448,7 @@ def freeze(
 
     See :meth:`ArrayContext.thaw`.
     """
-    if is_array_container(ary):
+    if is_array_container_type(ary.__class__):
         return map_array_container(partial(freeze, actx=actx), ary)
     else:
         if actx is None:
@@ -504,7 +504,7 @@ def from_numpy(ary: Any, actx: ArrayContext) -> Any:
     def _from_numpy(subary: Any) -> Any:
         if isinstance(subary, np.ndarray) and subary.dtype != "O":
             return actx.from_numpy(subary)
-        elif is_array_container(subary):
+        elif is_array_container_type(subary.__class__):
             return map_array_container(_from_numpy, subary)
         else:
             raise TypeError(f"unrecognized array type: '{type(subary).__name__}'")

arraycontext/fake_numpy.py

@@ -24,7 +24,7 @@
 
 
 import numpy as np
-from arraycontext.container import is_array_container, serialize_container
+from arraycontext.container import is_array_container_type, serialize_container
 from arraycontext.container.traversal import (
         rec_map_array_container, multimapped_over_array_containers)
 from pytools import memoize_in
@@ -182,7 +182,7 @@ def _new_like(self, ary, alloc_like):
             # e.g. `np.zeros_like(x)` returns `array([0, 0, ...], dtype=object)`
             # FIXME: what about object arrays nested in an ArrayContainer?
             raise NotImplementedError("operation not implemented for object arrays")
-        elif is_array_container(ary):
+        elif is_array_container_type(ary.__class__):
             return rec_map_array_container(alloc_like, ary)
         elif isinstance(ary, Number):
             # NOTE: `np.zeros_like(x)` returns `array(x, shape=())`, which
@@ -225,6 +225,26 @@ def _scalar_list_norm(ary, ord):
         raise NotImplementedError(f"unsupported value of 'ord': {ord}")
 
 
+def _reduce_norm(actx, arys, ord):
+    from numbers import Number
+    from functools import reduce
+
+    if ord is None:
+        ord = 2
+
+    # NOTE: these are ordered by an expected usage frequency
+    if ord == 2:
+        return actx.np.sqrt(sum(subary*subary for subary in arys))
+    elif ord == np.inf:
+        return reduce(actx.np.maximum, arys)
+    elif ord == -np.inf:
+        return reduce(actx.np.minimum, arys)
+    elif isinstance(ord, Number) and ord > 0:
+        return sum(subary**ord for subary in arys)**(1/ord)
+    else:
+        raise NotImplementedError(f"unsupported value of 'ord': {ord}")
+
+
 class BaseFakeNumpyLinalgNamespace:
     def __init__(self, array_context):
         self._array_context = array_context
@@ -253,8 +273,8 @@ def norm(self, ary, ord=None):
 
                 return flat_norm(ary, ord=ord)
 
-        if is_array_container(ary):
-            return _scalar_list_norm([
+        if is_array_container_type(ary.__class__):
+            return _reduce_norm(actx, [
                 self.norm(subary, ord=ord)
                 for _, subary in serialize_container(ary)
                 ], ord=ord)
@@ -266,14 +286,16 @@ def norm(self, ary, ord=None):
             raise NotImplementedError("only vector norms are implemented")
 
         if ary.size == 0:
-            return 0
+            return ary.dtype.type(0)
 
+        if ord == 2:
+            return actx.np.sqrt(actx.np.sum(abs(ary)**2))
         if ord == np.inf:
-            return self._array_context.np.max(abs(ary))
+            return actx.np.max(abs(ary))
         elif ord == -np.inf:
-            return self._array_context.np.min(abs(ary))
+            return actx.np.min(abs(ary))
         elif isinstance(ord, Number) and ord > 0:
-            return self._array_context.np.sum(abs(ary)**ord)**(1/ord)
+            return actx.np.sum(abs(ary)**ord)**(1/ord)
         else:
             raise NotImplementedError(f"unsupported value of 'ord': {ord}")
 # }}}

arraycontext/impl/pyopencl/__init__.py

@@ -244,11 +244,17 @@ def transform_loopy_program(self, t_unit):
 
             if "idof" in all_inames:
                 inner_iname = "idof"
+
         elif "i0" in all_inames:
             outer_iname = "i0"
 
             if "i1" in all_inames:
                 inner_iname = "i1"
+
+        elif not all_inames:
+            # no loops, nothing to transform
+            return t_unit
+
         else:
             raise RuntimeError(
                 "Unable to reason what outer_iname and inner_iname "

arraycontext/impl/pyopencl/fake_numpy.py

@@ -31,7 +31,7 @@
 
 from arraycontext.fake_numpy import \
         BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace
-from arraycontext.container import is_array_container
+from arraycontext.container import is_array_container_type
 from arraycontext.container.traversal import (
         rec_map_array_container,
         rec_multimap_array_container,
@@ -247,7 +247,7 @@ def as_device_scalar(bool_value):
         def rec_equal(x, y):
             if type(x) != type(y):
                 return as_device_scalar(False)
-            elif not is_array_container(x):
+            elif not is_array_container_type(x.__class__):
                 if x.shape != y.shape:
                     return as_device_scalar(False)
                 else:

arraycontext/impl/pytato/compile.py

@@ -30,7 +30,7 @@
 from arraycontext.container import ArrayContainer
 from arraycontext import PytatoPyOpenCLArrayContext
 from arraycontext.container.traversal import (rec_keyed_map_array_container,
-                                              is_array_container)
+                                              is_array_container_type)
 
 import numpy as np
 from typing import Any, Callable, Tuple, Dict, Mapping
@@ -124,7 +124,7 @@ def _get_arg_id_to_arg_and_arg_id_to_descr(args: Tuple[Any, ...],
             arg_id = (kw,)
             arg_id_to_arg[arg_id] = arg
             arg_id_to_descr[arg_id] = ScalarInputDescriptor(np.dtype(type(arg)))
-        elif is_array_container(arg):
+        elif is_array_container_type(arg.__class__):
             def id_collector(keys, ary):
                 arg_id = (kw,) + keys
                 arg_id_to_arg[arg_id] = ary
@@ -150,7 +150,7 @@ def _get_f_placeholder_args(arg, kw, arg_id_to_name):
     if np.isscalar(arg):
         name = arg_id_to_name[(kw,)]
         return pt.make_placeholder(name, (), np.dtype(type(arg)))
-    elif is_array_container(arg):
+    elif is_array_container_type(arg.__class__):
         def _rec_to_placeholder(keys, ary):
             name = arg_id_to_name[(kw,) + keys]
             return pt.make_placeholder(name, ary.shape, ary.dtype)
@@ -217,7 +217,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
                          **{kw: _get_f_placeholder_args(arg, kw, input_naming_map)
                             for kw, arg in kwargs.items()})
 
-        if not is_array_container(outputs):
+        if not is_array_container_type(outputs.__class__):
             # TODO: We could possibly just short-circuit this interface if the
             # returned type is a scalar. Not sure if it's worth it though.
             raise NotImplementedError(

arraycontext/impl/pytato/fake_numpy.py

@@ -26,7 +26,7 @@
 from arraycontext.fake_numpy import (
         BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace,
         )
-from arraycontext.container import is_array_container
+from arraycontext.container import is_array_container_type
 from arraycontext.container.traversal import (
         rec_map_array_container,
         rec_multimap_array_container,
@@ -179,7 +179,7 @@ def as_device_scalar(bool_value):
 
         if type(a) != type(b):
             return as_device_scalar(False)
-        elif not is_array_container(a):
+        elif not is_array_container_type(a.__class__):
             if a.shape != b.shape:
                 return as_device_scalar(False)
             else:

test/test_arraycontext.py

@@ -904,7 +904,7 @@ def test_numpy_conversion(actx_factory):
     with pytest.raises(TypeError):
         from_numpy(ac_actx, actx)
 
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         to_numpy(ac, actx)
 
 # }}}