Support containers in to/from_numpy

arraycontext/container/traversal.py

@@ -881,37 +881,24 @@ def from_numpy(
 
     The conversion is done using :meth:`arraycontext.ArrayContext.from_numpy`.
     """
-    def _from_numpy_with_check(subary: Union[np.ndarray, ScalarLike]) \
-            -> ArrayOrContainerOrScalar:
-        if isinstance(subary, np.ndarray) or np.isscalar(subary):
-            return actx.from_numpy(subary)
-        else:
-            raise TypeError(f"array is not an ndarray: '{type(subary).__name__}'")
+    warn("Calling from_numpy(ary, actx) is deprecated, call actx.from_numpy(ary)"
+         " instead. This will stop working in 2023.",
+         DeprecationWarning, stacklevel=2)
 
-    return rec_map_array_container(_from_numpy_with_check, ary)
+    return actx.from_numpy(ary)
 
 
-def to_numpy(ary: ArrayOrContainer, actx: ArrayContext) -> Any:
+def to_numpy(ary: ArrayOrContainer, actx: ArrayContext) -> ArrayOrContainer:
     """Convert all arrays in the :class:`~arraycontext.ArrayContainer` to
     :mod:`numpy` using the provided :class:`~arraycontext.ArrayContext` *actx*.
 
     The conversion is done using :meth:`arraycontext.ArrayContext.to_numpy`.
     """
-    def _to_numpy_with_check(subary: Any) -> Any:
-        if isinstance(subary, actx.array_types) or np.isscalar(subary):
-            # NOTE: these are allowed by np.isscalar, but not here
-            assert not isinstance(subary, (str, bytes))
-
-            return actx.to_numpy(subary)
-        else:
-            raise TypeError(
-                    f"array of type '{type(subary).__name__}' not in "
-                    f"supported types {actx.array_types}")
+    warn("Calling to_numpy(ary, actx) is deprecated, call actx.to_numpy(ary)"
+         " instead. This will stop working in 2023.",
+         DeprecationWarning, stacklevel=2)
 
-    return rec_map_array_container(_to_numpy_with_check,
-                                   # do a freeze first, if 'actx' supports
-                                   # container-wide freezes
-                                   actx.thaw(actx.freeze(ary)))
+    return actx.to_numpy(ary)
 
 # }}}
 

arraycontext/context.py

@@ -221,6 +221,8 @@ def dtype(self) -> "np.dtype[Any]":
         "ArrayOrContainerOrScalarT",
         bound=ArrayOrContainerOrScalar)
 
+NumpyOrContainerOrScalar = Union[np.ndarray, "ArrayContainer", ScalarLike]
+
 # }}}
 
 
@@ -304,25 +306,28 @@ def zeros_like(self, ary: Array) -> Array:
 
     @abstractmethod
     def from_numpy(self,
-                   array: Union["np.ndarray[Any, Any]", ScalarLike]
-                   ) -> Union[Array, ScalarLike]:
+                   array: ArrayOrContainerOrScalar
+                   ) -> NumpyOrContainerOrScalar:
         r"""
         :returns: the :class:`numpy.ndarray` *array* converted to the
             array context's array type. The returned array will be
-            :meth:`thaw`\ ed.
+            :meth:`thaw`\ ed. When working with array containers each leaf
+            must be an :class:`~numpy.ndarray` or scalar, which is then converted
+            to the context's array type leaving the container structure
+            intact.
         """
-        pass
 
     @abstractmethod
     def to_numpy(self,
-                 array: Union[Array, ScalarLike]
-                 ) -> Union["np.ndarray[Any, Any]", ScalarLike]:
+                 array: NumpyOrContainerOrScalar
+                 ) -> ArrayOrContainerOrScalar:
         r"""
-        :returns: *array*, an array recognized by the context, converted
-            to a :class:`numpy.ndarray`. *array* must be
-            :meth:`thaw`\ ed.
+        :returns: an :class:`numpy.ndarray` for each array recognized by the
+            context. The input *array* must be :meth:`thaw`\ ed.
+            When working with array containers each leaf must be one of
+            the context's array types or a scalar, which is then converted to
+            an :class:`~numpy.ndarray` leaving the container structure intact.
         """
-        pass
 
     @abstractmethod
     def call_loopy(self,
@@ -351,8 +356,6 @@ def freeze(self, array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
         Freezing makes the array independent of this :class:`ArrayContext`;
         it is permitted to :meth:`thaw` it in a different one, as long as that
         context understands the array format.
-
-        See also :func:`arraycontext.freeze`.
         """
 
     @abstractmethod
@@ -366,8 +369,6 @@ def thaw(self, array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
         the data in *array*.
 
         The returned array may not be used with other contexts while thawed.
-
-        See also :func:`arraycontext.thaw`.
         """
 
     def freeze_thaw(
@@ -387,10 +388,11 @@ def freeze_thaw(
     @abstractmethod
     def tag(self,
             tags: ToTagSetConvertible,
-            array: ArrayT) -> ArrayT:
+            array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
         """If the array type used by the array context is capable of capturing
         metadata, return a version of *array* with the *tags* applied. *array*
-        itself is not modified.
+        itself is not modified. When working with array containers, the
+        tags are applied to each leaf of the container.
 
         See :ref:`metadata` as well as application-specific metadata types.
 
@@ -400,10 +402,11 @@ def tag(self,
     @abstractmethod
     def tag_axis(self,
                  iaxis: int, tags: ToTagSetConvertible,
-                 array: ArrayT) -> ArrayT:
+                 array: ArrayOrContainerOrScalarT) -> ArrayOrContainerOrScalarT:
         """If the array type used by the array context is capable of capturing
         metadata, return a version of *array* in which axis number *iaxis* has
-        the *tags* applied. *array* itself is not modified.
+        the *tags* applied. *array* itself is not modified. When working with
+        array containers, the tags are applied to each leaf of the container.
 
         See :ref:`metadata` as well as application-specific metadata types.
 

arraycontext/fake_numpy.py

@@ -91,25 +91,11 @@ def _get_fake_numpy_linalg_namespace(self):
         # "interp",
         })
 
-    def _new_like(self, ary, alloc_like):
-        if np.isscalar(ary):
-            # NOTE: `np.zeros_like(x)` returns `array(x, shape=())`, which
-            # is best implemented by concrete array contexts, if at all
-            raise NotImplementedError("operation not implemented for scalars")
-
-        if isinstance(ary, np.ndarray) and ary.dtype.char == "O":
-            # NOTE: we don't want to match numpy semantics on object arrays,
-            # 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")
-
-        return rec_map_array_container(alloc_like, ary)
-
     def empty_like(self, ary):
-        return self._new_like(ary, self._array_context.empty_like)
+        return self._array_context.empty_like(ary)
 
     def zeros_like(self, ary):
-        return self._new_like(ary, self._array_context.zeros_like)
+        return self._array_context.zeros_like(ary)
 
     def conjugate(self, x):
         # NOTE: conjugate distributes over object arrays, but it looks for a

arraycontext/impl/jax/__init__.py

@@ -27,11 +27,12 @@
 THE SOFTWARE.
 """
 
+from typing import Callable, Optional, Tuple
+
 import numpy as np
 
-from typing import Union, Callable, Any
 from pytools.tag import ToTagSetConvertible
-from arraycontext.context import ArrayContext, ScalarLike
+from arraycontext.context import ArrayContext, Array, ArrayOrContainer, ScalarLike
 from arraycontext.container.traversal import (with_array_context,
                                               rec_map_array_container)
 
@@ -60,6 +61,30 @@ def _get_fake_numpy_namespace(self):
         from .fake_numpy import EagerJAXFakeNumpyNamespace
         return EagerJAXFakeNumpyNamespace(self)
 
+    def _rec_map_container(
+            self, func: Callable[[Array], Array], array: ArrayOrContainer,
+            allowed_types: Optional[Tuple[type, ...]] = None, *,
+            default_scalar: Optional[ScalarLike] = None,
+            strict: bool = False) -> ArrayOrContainer:
+        if allowed_types is None:
+            allowed_types = self.array_types
+
+        def _wrapper(ary):
+            if isinstance(ary, allowed_types):
+                return func(ary)
+            elif np.isscalar(ary):
+                if default_scalar is None:
+                    return ary
+                else:
+                    return np.array(ary).dtype.type(default_scalar)
+            else:
+                raise TypeError(
+                    f"{type(self).__name__}.{func.__name__[1:]} invoked with "
+                    f"an unsupported array type: got '{type(ary).__name__}', "
+                    f"but expected one of {allowed_types}")
+
+        return rec_map_array_container(_wrapper, array)
+
     # {{{ ArrayContext interface
 
     def empty(self, shape, dtype):
@@ -70,40 +95,45 @@ def zeros(self, shape, dtype):
         import jax.numpy as jnp
         return jnp.zeros(shape=shape, dtype=dtype)
 
-    def from_numpy(self, array: Union[np.ndarray, ScalarLike]):
-        import jax
-        return jax.device_put(array)
+    def empty_like(self, ary):
+        def _empty_like(array):
+            return self.empty(array.shape, array.dtype)
+
+        return self._rec_map_container(_empty_like, ary)
+
+    def zeros_like(self, ary):
+        def _zeros_like(array):
+            return self.zeros(array.shape, array.dtype)
+
+        return self._rec_map_container(_zeros_like, ary, default_scalar=0)
+
+    def from_numpy(self, array):
+        def _from_numpy(ary):
+            import jax
+            return jax.device_put(ary)
+
+        return with_array_context(
+            self._rec_map_container(_from_numpy, array, allowed_types=(np.ndarray,)),
+            actx=self)
 
     def to_numpy(self, array):
-        import jax
-        # jax.device_get can take scalars as well.
-        return jax.device_get(array)
+        def _to_numpy(ary):
+            import jax
+            return jax.device_get(ary)
 
-    def call_loopy(self, t_unit, **kwargs):
-        raise NotImplementedError("calling loopy on JAX arrays"
-                                  " not supported. Maybe rewrite"
-                                  " the loopy kernel as numpy-flavored array"
-                                  " operations using ArrayContext.np.")
+        return with_array_context(
+            self._rec_map_container(_to_numpy, array),
+            actx=None)
 
     def freeze(self, array):
-        from jax.numpy import DeviceArray
-
-        def _rec_freeze(ary):
-            if isinstance(ary, DeviceArray):
-                pass
-            else:
-                raise TypeError(f"{type(self).__name__}.thaw expects "
-                                f"`jax.DeviceArray` got {type(ary)}.")
+        def _freeze(ary):
             return ary.block_until_ready()
 
-        return with_array_context(rec_map_array_container(_rec_freeze, array),
-                                  actx=None)
+        return with_array_context(self._rec_map_container(_freeze, array), actx=None)
 
     def thaw(self, array):
         return with_array_context(array, actx=self)
 
-    # }}}
-
     def tag(self, tags: ToTagSetConvertible, array):
         # Sorry, not capable.
         return array
@@ -112,21 +142,28 @@ def tag_axis(self, iaxis, tags: ToTagSetConvertible, array):
         # TODO: See `jax.experiemental.maps.xmap`, proabably that should be useful?
         return array
 
-    def clone(self):
-        return type(self)()
-
-    def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
-        return f
+    def call_loopy(self, t_unit, **kwargs):
+        raise NotImplementedError(
+            "Calling loopy on JAX arrays is not supported. Maybe rewrite"
+            " the loopy kernel as numpy-flavored array operations using"
+            " ArrayContext.np.")
 
     def einsum(self, spec, *args, arg_names=None, tagged=()):
         import jax.numpy as jnp
         if arg_names is not None:
             from warnings import warn
             warn("'arg_names' don't bear any significance in "
-                 "EagerJAXArrayContext.", stacklevel=2)
+                 f"{type(self).__name__}.", stacklevel=2)
 
         return jnp.einsum(spec, *args)
 
+    def clone(self):
+        return type(self)()
+
+    # }}}
+
+    # {{{ properties
+
     @property
     def permits_inplace_modification(self):
         return False
@@ -139,4 +176,6 @@ def supports_nonscalar_broadcasting(self):
     def permits_advanced_indexing(self):
         return True
 
+    # }}}
+
 # vim: foldmethod=marker

arraycontext/impl/jax/fake_numpy.py

@@ -23,6 +23,9 @@
 """
 from functools import partial, reduce
 
+import numpy as np
+import jax.numpy as jnp
+
 from arraycontext.fake_numpy import (
         BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace,
         )
@@ -31,8 +34,6 @@
         rec_map_reduce_array_container,
         )
 from arraycontext.container import NotAnArrayContainerError, serialize_container
-import numpy
-import jax.numpy as jnp
 
 
 class EagerJAXFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
@@ -60,9 +61,10 @@ def ones_like(self, ary):
 
     def full_like(self, ary, fill_value):
         def _full_like(subary):
-            return jnp.full_like(ary, fill_value)
+            return jnp.full_like(subary, fill_value)
 
-        return self._new_like(ary, _full_like)
+        return self._array_context._rec_map_container(
+            _full_like, ary, default_scalar=fill_value)
 
     # }}}
 
@@ -111,11 +113,10 @@ def vdot(self, x, y, dtype=None):
         from arraycontext import rec_multimap_reduce_array_container
 
         def _rec_vdot(ary1, ary2):
-            if dtype not in [None, numpy.find_common_type((ary1.dtype,
-                                                           ary2.dtype),
-                                                          ())]:
-                raise NotImplementedError(f"{type(self)} cannot take dtype in"
-                                          " vdot.")
+            common_dtype = np.find_common_type((ary1.dtype, ary2.dtype), ())
+            if dtype not in [None, common_dtype]:
+                raise NotImplementedError(
+                    f"{type(self).__name__} cannot take dtype in vdot.")
 
             return jnp.vdot(ary1, ary2)
 
@@ -125,12 +126,20 @@ def _rec_vdot(ary1, ary2):
 
     # {{{ logic functions
 
+    def all(self, a):
+        return rec_map_reduce_array_container(
+            partial(reduce, jnp.logical_and), jnp.all, a)
+
+    def any(self, a):
+        return rec_map_reduce_array_container(
+            partial(reduce, jnp.logical_or), jnp.any, a)
+
     def array_equal(self, a, b):
         actx = self._array_context
 
         # NOTE: not all backends support `bool` properly, so use `int8` instead
-        true = actx.from_numpy(numpy.int8(True))
-        false = actx.from_numpy(numpy.int8(False))
+        true = actx.from_numpy(np.int8(True))
+        false = actx.from_numpy(np.int8(False))
 
         def rec_equal(x, y):
             if type(x) != type(y):