diff --git a/arraycontext/container/traversal.py b/arraycontext/container/traversal.py index 6c858fe3..642aaf16 100644 --- a/arraycontext/container/traversal.py +++ b/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) # }}} diff --git a/arraycontext/context.py b/arraycontext/context.py index 727e1ba9..36a7acee 100644 --- a/arraycontext/context.py +++ b/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. diff --git a/arraycontext/fake_numpy.py b/arraycontext/fake_numpy.py index 73c9e406..d5c8fce9 100644 --- a/arraycontext/fake_numpy.py +++ b/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 diff --git a/arraycontext/impl/jax/__init__.py b/arraycontext/impl/jax/__init__.py index 81076ae0..dfb89c45 100644 --- a/arraycontext/impl/jax/__init__.py +++ b/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 diff --git a/arraycontext/impl/jax/fake_numpy.py b/arraycontext/impl/jax/fake_numpy.py index 394fbd52..37c99b4a 100644 --- a/arraycontext/impl/jax/fake_numpy.py +++ b/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): diff --git a/arraycontext/impl/pyopencl/__init__.py b/arraycontext/impl/pyopencl/__init__.py index 14672465..9ad36251 100644 --- a/arraycontext/impl/pyopencl/__init__.py +++ b/arraycontext/impl/pyopencl/__init__.py @@ -29,13 +29,13 @@ """ from warnings import warn -from typing import Dict, List, Optional, Union, TYPE_CHECKING +from typing import Callable, Dict, List, Optional, Tuple, TYPE_CHECKING import numpy as np from pytools.tag import ToTagSetConvertible -from arraycontext.context import ArrayContext, ScalarLike +from arraycontext.context import ArrayContext, Array, ArrayOrContainer, ScalarLike from arraycontext.container.traversal import (rec_map_array_container, with_array_context) @@ -151,36 +151,111 @@ def _get_fake_numpy_namespace(self): from arraycontext.impl.pyopencl.fake_numpy import PyOpenCLFakeNumpyNamespace return PyOpenCLFakeNumpyNamespace(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: + import arraycontext.impl.pyopencl.taggable_cl_array as tga + + if allowed_types is None: + # TODO: replace with 'self.array_types' once `cla.Array` support + # is completely removed + allowed_types = (tga.TaggableCLArray,) + + def _wrapper(ary): + if isinstance(ary, allowed_types): + return func(ary) + elif not strict and isinstance(ary, self.array_types): + from warnings import warn + warn(f"Invoking {type(self).__name__}.{func.__name__[1:]} with " + f"{type(ary).__name__} will be unsupported in 2023. Use " + "'to_tagged_cl_array' to convert instances to TaggableCLArray.", + DeprecationWarning, stacklevel=2) + return func(tga.to_tagged_cl_array(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): - from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray - - return TaggableCLArray(self.queue, shape=shape, dtype=dtype, - allocator=self.allocator) + import arraycontext.impl.pyopencl.taggable_cl_array as tga + return tga.empty(self.queue, shape, dtype, allocator=self.allocator) def zeros(self, shape, dtype): - import pyopencl.array as cl_array - from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array - return to_tagged_cl_array(cl_array.zeros(self.queue, shape=shape, - dtype=dtype, - allocator=self.allocator), - axes=None, tags=frozenset()) + import arraycontext.impl.pyopencl.taggable_cl_array as tga + return tga.zeros(self.queue, shape, dtype, allocator=self.allocator) - def from_numpy(self, array: Union[np.ndarray, ScalarLike]): - import pyopencl.array as cl_array - from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array - return to_tagged_cl_array(cl_array - .to_device(self.queue, - array, - allocator=self.allocator), - axes=None, tags=frozenset()) + def empty_like(self, ary): + import arraycontext.impl.pyopencl.taggable_cl_array as tga + + def _empty_like(array): + return tga.empty(self.queue, array.shape, array.dtype, + allocator=self.allocator, axes=array.axes, tags=array.tags) + + return self._rec_map_container(_empty_like, ary) + + def zeros_like(self, ary): + import arraycontext.impl.pyopencl.taggable_cl_array as tga + + def _zeros_like(array): + return tga.zeros(self.queue, array.shape, array.dtype, + allocator=self.allocator, axes=array.axes, tags=array.tags) + + return self._rec_map_container(_zeros_like, ary, default_scalar=0) + + def from_numpy(self, array): + import arraycontext.impl.pyopencl.taggable_cl_array as tga + + def _from_numpy(ary): + return tga.to_device(self.queue, ary, allocator=self.allocator) + + return with_array_context( + self._rec_map_container(_from_numpy, array, (np.ndarray,), strict=True), + actx=self) def to_numpy(self, array): - if np.isscalar(array): - return array + def _to_numpy(ary): + return ary.get(queue=self.queue) + + return with_array_context( + self._rec_map_container(_to_numpy, array), + actx=None) + + def freeze(self, array): + def _freeze(ary): + ary.finish() + return ary.with_queue(None) - return array.get(queue=self.queue) + return with_array_context(self._rec_map_container(_freeze, array), actx=None) + + def thaw(self, array): + def _thaw(ary): + return ary.with_queue(self.queue) + + return with_array_context(self._rec_map_container(_thaw, array), actx=self) + + def tag(self, tags: ToTagSetConvertible, array): + def _tag(ary): + return ary.tagged(tags) + + return self._rec_map_container(_tag, array) + + def tag_axis(self, iaxis: int, tags: ToTagSetConvertible, array): + def _tag_axis(ary): + return ary.with_tagged_axis(iaxis, tags) + + return self._rec_map_container(_tag_axis, array) def call_loopy(self, t_unit, **kwargs): try: @@ -202,50 +277,19 @@ def call_loopy(self, t_unit, **kwargs): if len(wait_event_queue) > self._wait_event_queue_length: wait_event_queue.pop(0).wait() - from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array + import arraycontext.impl.pyopencl.taggable_cl_array as tga # FIXME: Inherit loopy tags for these arrays - return {name: to_tagged_cl_array(ary, axes=None, tags=frozenset()) - for name, ary in result.items()} - - def freeze(self, array): - import pyopencl.array as cl_array + return {name: tga.to_tagged_cl_array(ary) for name, ary in result.items()} - def _rec_freeze(ary): - if isinstance(ary, cl_array.Array): - ary.finish() - return ary.with_queue(None) - else: - raise TypeError(f"{type(self).__name__} cannot freeze" - f" arrays of type '{type(ary).__name__}'.") - - return with_array_context(rec_map_array_container(_rec_freeze, array), - actx=None) - - def thaw(self, array): - from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, - to_tagged_cl_array) - import pyopencl.array as cl_array - - def _rec_thaw(ary): - if isinstance(ary, TaggableCLArray): - return ary.with_queue(self.queue) - elif isinstance(ary, cl_array.Array): - from warnings import warn - warn("Invoking PyOpenCLArrayContext.thaw with pyopencl.Array" - " will be unsupported in 2023. Use `to_tagged_cl_array`" - " to convert instances of pyopencl.Array to TaggableCLArray.", - DeprecationWarning, stacklevel=2) - return (to_tagged_cl_array(ary, axes=None, tags=frozenset()) - .with_queue(self.queue)) - else: - raise ValueError("array should be a cl.array.Array," - f" got '{type(ary)}'") - - return with_array_context(rec_map_array_container(_rec_thaw, array), - actx=self) + def clone(self): + return type(self)(self.queue, self.allocator, + wait_event_queue_length=self._wait_event_queue_length, + force_device_scalars=self._force_device_scalars) # }}} + # {{{ transform_loopy_program + def transform_loopy_program(self, t_unit): from warnings import warn warn("Using arraycontext.PyOpenCLArrayContext.transform_loopy_program " @@ -316,43 +360,9 @@ def transform_loopy_program(self, t_unit): return t_unit - def tag(self, tags: ToTagSetConvertible, array): - import pyopencl.array as cl_array - from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, - to_tagged_cl_array) - - def _rec_tagged(ary): - if isinstance(ary, TaggableCLArray): - return ary.tagged(tags) - elif isinstance(ary, cl_array.Array): - return to_tagged_cl_array(ary, axes=None, tags=tags) - else: - raise ValueError("array should be a cl.array.Array," - f" got '{type(ary)}'") - - return rec_map_array_container(_rec_tagged, array) - - def tag_axis(self, iaxis, tags: ToTagSetConvertible, array): - import pyopencl.array as cl_array - from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, - to_tagged_cl_array) - - def _rec_tagged(ary): - if isinstance(ary, TaggableCLArray): - return ary.with_tagged_axis(iaxis, tags) - elif isinstance(ary, cl_array.Array): - return (to_tagged_cl_array(ary, axes=None, tags=tags) - .with_tagged_axis(iaxis, tags)) - else: - raise ValueError("array should be a cl.array.Array," - f" got '{type(ary)}'") - - return rec_map_array_container(_rec_tagged, array) + # }}} - def clone(self): - return type(self)(self.queue, self.allocator, - wait_event_queue_length=self._wait_event_queue_length, - force_device_scalars=self._force_device_scalars) + # {{{ properties @property def permits_inplace_modification(self): @@ -366,6 +376,8 @@ def supports_nonscalar_broadcasting(self): def permits_advanced_indexing(self): return False + # }}} + # }}} # vim: foldmethod=marker diff --git a/arraycontext/impl/pyopencl/fake_numpy.py b/arraycontext/impl/pyopencl/fake_numpy.py index 3c9be87d..2e206a8b 100644 --- a/arraycontext/impl/pyopencl/fake_numpy.py +++ b/arraycontext/impl/pyopencl/fake_numpy.py @@ -67,18 +67,24 @@ def ones_like(self, ary): return self.full_like(ary, 1) def full_like(self, ary, fill_value): + import arraycontext.impl.pyopencl.taggable_cl_array as tga + def _full_like(subary): - ones = self._array_context.empty_like(subary) - ones.fill(fill_value) - return ones + filled = tga.empty( + self._array_context.queue, subary.shape, subary.dtype, + allocator=self._array_context.allocator, + axes=subary.axes, tags=subary.tags) + filled.fill(fill_value) + return filled - return self._new_like(ary, _full_like) + return self._array_context._rec_map_container( + _full_like, ary, default_scalar=fill_value) def copy(self, ary): def _copy(subary): return subary.copy(queue=self._array_context.queue) - return self._new_like(ary, _copy) + return self._array_context._rec_map_container(_copy, ary) # }}} @@ -144,9 +150,15 @@ def vdot(self, x, y, dtype=None): def all(self, a): queue = self._array_context.queue + + def _all(ary): + if np.isscalar(ary): + return np.int8(all([ary])) + return ary.all(queue=queue) + result = rec_map_reduce_array_container( partial(reduce, partial(cl_array.minimum, queue=queue)), - lambda subary: subary.all(queue=queue), + _all, a) if not self._array_context._force_device_scalars: @@ -155,9 +167,15 @@ def all(self, a): def any(self, a): queue = self._array_context.queue + + def _any(ary): + if np.isscalar(ary): + return np.int8(any([ary])) + return ary.any(queue=queue) + result = rec_map_reduce_array_container( partial(reduce, partial(cl_array.maximum, queue=queue)), - lambda subary: subary.any(queue=queue), + _any, a) if not self._array_context._force_device_scalars: diff --git a/arraycontext/impl/pyopencl/taggable_cl_array.py b/arraycontext/impl/pyopencl/taggable_cl_array.py index 439ca580..49ae08b1 100644 --- a/arraycontext/impl/pyopencl/taggable_cl_array.py +++ b/arraycontext/impl/pyopencl/taggable_cl_array.py @@ -5,12 +5,17 @@ .. autofunction:: to_tagged_cl_array """ -import pyopencl.array as cla -from typing import Any, Dict, FrozenSet, Optional, Tuple -from pytools.tag import Taggable, Tag, ToTagSetConvertible from dataclasses import dataclass +from typing import Any, Dict, FrozenSet, Optional, Tuple + +import numpy as np +import pyopencl.array as cla + from pytools import memoize +from pytools.tag import Taggable, Tag, ToTagSetConvertible + +# {{{ utils @dataclass(frozen=True, eq=True) class Axis(Taggable): @@ -42,6 +47,10 @@ def _unwrap_cl_array(ary: cla.Array) -> Dict[str, Any]: _fast=True, ) +# }}} + + +# {{{ TaggableCLArray class TaggableCLArray(cla.Array, Taggable): """ @@ -111,8 +120,8 @@ def with_tagged_axis(self, iaxis: int, def to_tagged_cl_array(ary: cla.Array, - axes: Optional[Tuple[Axis, ...]], - tags: FrozenSet[Tag]) -> TaggableCLArray: + axes: Optional[Tuple[Axis, ...]] = None, + tags: FrozenSet[Tag] = frozenset()) -> TaggableCLArray: """ Returns a :class:`TaggableCLArray` that is constructed from the data in *ary* along with the metadata from *axes* and *tags*. If *ary* is already a @@ -123,7 +132,7 @@ def to_tagged_cl_array(ary: cla.Array, array. If passed *None*, then initialized to a :class:`pytato.Axis` with no tags attached for each dimension. """ - if axes and len(axes) != ary.ndim: + if axes is not None and len(axes) != ary.ndim: raise ValueError("axes length does not match array dimension: " f"got {len(axes)} axes for {ary.ndim}d array") @@ -131,7 +140,7 @@ def to_tagged_cl_array(ary: cla.Array, tags = normalize_tags(tags) if isinstance(ary, TaggableCLArray): - if axes: + if axes is not None: for i, axis in enumerate(axes): ary = ary.with_tagged_axis(i, axis.tags) @@ -144,3 +153,45 @@ def to_tagged_cl_array(ary: cla.Array, **_unwrap_cl_array(ary)) else: raise TypeError(f"unsupported array type: '{type(ary).__name__}'") + +# }}} + + +# {{{ creation + +def empty(queue, shape, dtype=float, *, + axes: Optional[Tuple[Axis, ...]] = None, + tags: FrozenSet[Tag] = frozenset(), + order: str = "C", + allocator=None) -> TaggableCLArray: + if dtype is not None: + dtype = np.dtype(dtype) + + return TaggableCLArray( + queue, shape, dtype, + axes=axes, tags=tags, + order=order, allocator=allocator) + + +def zeros(queue, shape, dtype=float, *, + axes: Optional[Tuple[Axis, ...]] = None, + tags: FrozenSet[Tag] = frozenset(), + order: str = "C", + allocator=None) -> TaggableCLArray: + result = empty( + queue, shape, dtype=dtype, axes=axes, tags=tags, + order=order, allocator=allocator) + result._zero_fill() + + return result + + +def to_device(queue, ary, *, + axes: Optional[Tuple[Axis, ...]] = None, + tags: FrozenSet[Tag] = frozenset(), + allocator=None): + return to_tagged_cl_array( + cla.to_device(queue, ary, allocator=allocator), + axes=axes, tags=tags) + +# }}} diff --git a/arraycontext/impl/pytato/__init__.py b/arraycontext/impl/pytato/__init__.py index ed3bff8d..8d7e0426 100644 --- a/arraycontext/impl/pytato/__init__.py +++ b/arraycontext/impl/pytato/__init__.py @@ -42,17 +42,18 @@ THE SOFTWARE. """ +import abc import sys -from arraycontext.context import ArrayContext, ScalarLike -from arraycontext.container.traversal import (rec_map_array_container, - with_array_context) -from arraycontext.metadata import NameHint +from typing import (Any, Callable, Union, Tuple, Type, FrozenSet, Dict, Optional, + TYPE_CHECKING) import numpy as np -from typing import (Any, Callable, Union, TYPE_CHECKING, Tuple, Type, FrozenSet, - Dict, Optional) from pytools.tag import ToTagSetConvertible, normalize_tags, Tag -import abc + +from arraycontext.context import ArrayContext, Array, ArrayOrContainer, ScalarLike +from arraycontext.container.traversal import (rec_map_array_container, + with_array_context) +from arraycontext.metadata import NameHint if TYPE_CHECKING: import pytato @@ -105,9 +106,11 @@ class _BasePytatoArrayContext(ArrayContext, abc.ABC): .. automethod:: compile """ - def __init__(self, - *, compile_trace_callback: Optional[Callable[[Any, str, Any], None]] - = None) -> None: + + def __init__( + self, *, + compile_trace_callback: Optional[Callable[[Any, str, Any], None]] = None + ) -> None: """ :arg compile_trace_callback: A function of three arguments *(what, stage, ir)*, where *what* identifies the object @@ -116,9 +119,10 @@ def __init__(self, representation. This interface should be considered unstable. """ + super().__init__() + import pytato as pt import loopy as lp - super().__init__() self._freeze_prg_cache: Dict[pt.DictOfNamedArrays, lp.TranslationUnit] = {} self._dag_transform_cache: Dict[ pt.DictOfNamedArrays, @@ -136,13 +140,30 @@ def _get_fake_numpy_namespace(self): from arraycontext.impl.pytato.fake_numpy import PytatoFakeNumpyNamespace return PytatoFakeNumpyNamespace(self) + @abc.abstractproperty + def _frozen_array_types(self) -> Tuple[Type, ...]: + """ + Returns valid frozen array types for the array context. + """ + + # {{{ ArrayContext interface + def empty(self, shape, dtype): - raise ValueError(f"{type(self).__name__} does not support empty") + raise NotImplementedError( + f"{type(self).__name__}.empty is not supported") def zeros(self, shape, dtype): import pytato as pt return pt.zeros(shape, dtype) + def empty_like(self, ary): + raise NotImplementedError( + f"{type(self).__name__}.empty_like is not supported") + + # }}} + + # {{{ compilation + def transform_dag(self, dag: "pytato.DictOfNamedArrays" ) -> "pytato.DictOfNamedArrays": """ @@ -158,21 +179,18 @@ def transform_dag(self, dag: "pytato.DictOfNamedArrays" return dag def transform_loopy_program(self, t_unit): - raise ValueError(f"{type(self)} does not implement " - "transform_loopy_program. Sub-classes are supposed " - "to implement it.") - - @abc.abstractproperty - def frozen_array_types(self) -> Tuple[Type, ...]: - """ - Returns valid frozen array types for the array context. - """ - pass + raise ValueError( + f"{type(self).__name__} does not implement transform_loopy_program. " + "Sub-classes are supposed to implement it.") @abc.abstractmethod def einsum(self, spec, *args, arg_names=None, tagged=()): pass + # }}} + + # {{{ properties + @property def permits_inplace_modification(self): return False @@ -185,6 +203,8 @@ def supports_nonscalar_broadcasting(self): def permits_advanced_indexing(self): return True + # }}} + # }}} @@ -210,10 +230,10 @@ class PytatoPyOpenCLArrayContext(_BasePytatoArrayContext): .. automethod:: compile """ - def __init__(self, queue: "cl.CommandQueue", allocator=None, - *, - compile_trace_callback: Optional[Callable[[Any, str, Any], None]] - = None) -> None: + def __init__( + self, queue: "cl.CommandQueue", allocator=None, *, + compile_trace_callback: Optional[Callable[[Any, str, Any], None]] = None + ) -> None: """ :arg compile_trace_callback: A function of three arguments *(what, stage, ir)*, where *what* identifies the object @@ -232,63 +252,82 @@ def __init__(self, queue: "cl.CommandQueue", allocator=None, # unused, but necessary to keep the context alive self.context = self.queue.context - # {{{ ArrayContext interface - - def clone(self): - return type(self)(self.queue, self.allocator) + @property + def _frozen_array_types(self) -> Tuple[Type, ...]: + import pyopencl.array as cla + return (cla.Array,) - def from_numpy(self, array: Union[np.ndarray, ScalarLike]): + 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: import pytato as pt - import pyopencl.array as cla - cl_array = cla.to_device(self.queue, array) - return pt.make_data_wrapper(cl_array) + import arraycontext.impl.pyopencl.taggable_cl_array as tga - def to_numpy(self, array): - if np.isscalar(array): - return array + if allowed_types is None: + allowed_types = (pt.Array, tga.TaggableCLArray) - cl_array = self.freeze(array) - return cl_array.get(queue=self.queue) + def _wrapper(ary): + if isinstance(ary, allowed_types): + return func(ary) + elif not strict and isinstance(ary, self._frozen_array_types): + from warnings import warn + warn(f"Invoking {type(self).__name__}.{func.__name__[1:]} with" + f" {type(ary).__name__} will be unsupported in 2023. Use" + " 'to_tagged_cl_array' to convert instances to" + " TaggableCLArray.", DeprecationWarning, stacklevel=2) + + return func(tga.to_tagged_cl_array(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}") - @property - def frozen_array_types(self) -> Tuple[Type, ...]: - import pyopencl.array as cla - return (cla.Array, ) + return rec_map_array_container(_wrapper, array) - def call_loopy(self, program, **kwargs): - import pytato as pt - from pytato.scalar_expr import SCALAR_CLASSES - from pytato.loopy import call_loopy - from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray + # {{{ ArrayContext interface - entrypoint = program.default_entrypoint.name + def zeros_like(self, ary): + def _zeros_like(array): + return self.zeros(array.shape, array.dtype) - # {{{ preprocess args + return self._rec_map_container(_zeros_like, ary, default_scalar=0) - processed_kwargs = {} + def from_numpy(self, array): + import pytato as pt + import arraycontext.impl.pyopencl.taggable_cl_array as tga - for kw, arg in sorted(kwargs.items()): - if isinstance(arg, (pt.Array,) + SCALAR_CLASSES): - pass - elif isinstance(arg, TaggableCLArray): - arg = self.thaw(arg) - else: - raise ValueError(f"call_loopy argument '{kw}' expected to be an" - " instance of 'pytato.Array', 'Number' or" - f"'TaggableCLArray', got '{type(arg)}'") + def _from_numpy(ary): + return pt.make_data_wrapper( + tga.to_device(self.queue, ary, allocator=self.allocator) + ) - processed_kwargs[kw] = arg + return with_array_context( + self._rec_map_container(_from_numpy, array, (np.ndarray,), strict=True), + actx=self) - # }}} + def to_numpy(self, array): + def _to_numpy(ary): + return ary.get(queue=self.queue) - return call_loopy(program, processed_kwargs, entrypoint) + return with_array_context( + self._rec_map_container(_to_numpy, self.freeze(array)), + actx=None) def freeze(self, array): + if np.isscalar(array): + return array + import pytato as pt import pyopencl.array as cla - import loopy as lp - from arraycontext.context import ArrayT from arraycontext.container.traversal import rec_keyed_map_array_container from arraycontext.impl.pytato.utils import (_normalize_pt_expr, get_cl_axes_from_pt_axes) @@ -296,16 +335,15 @@ def freeze(self, array): TaggableCLArray) from arraycontext.impl.pytato.compile import _ary_container_key_stringifier - array_as_dict: Dict[str, Union[cla.Array, TaggableCLArray, - pt.Array]] = {} + array_as_dict: Dict[str, Union[cla.Array, TaggableCLArray, pt.Array]] = {} key_to_frozen_subary: Dict[str, TaggableCLArray] = {} key_to_pt_arrays: Dict[str, pt.Array] = {} - def _record_leaf_ary_in_dict(key: Tuple[Any, ...], - ary: ArrayT): + def _record_leaf_ary_in_dict( + key: Tuple[Any, ...], + ary: Union[cla.Array, TaggableCLArray, pt.Array]) -> None: key_str = "_ary" + _ary_container_key_stringifier(key) array_as_dict[key_str] = ary - return ary rec_keyed_map_array_container(_record_leaf_ary_in_dict, array) @@ -314,37 +352,37 @@ def _record_leaf_ary_in_dict(key: Tuple[Any, ...], for key, subary in array_as_dict.items(): if isinstance(subary, TaggableCLArray): key_to_frozen_subary[key] = subary.with_queue(None) - elif isinstance(subary, cla.Array): + elif isinstance(subary, self._frozen_array_types): from warnings import warn - warn("Freezing pyopencl.array.Array will be deprecated in 2023." - " Use `to_tagged_cl_array` to convert the array to" - " TaggableCLArray", DeprecationWarning, stacklevel=2) - key_to_frozen_subary[key] = to_tagged_cl_array( - subary.with_queue(None), - axes=None, - tags=frozenset()) + warn(f"Invoking {type(self).__name__}.freeze with" + f" {type(subary).__name__} will be unsupported in 2023. Use" + " `to_tagged_cl_array` to convert instances to TaggableCLArray.", + DeprecationWarning, stacklevel=2) + + key_to_frozen_subary[key] = ( + to_tagged_cl_array(subary.with_queue(None))) elif isinstance(subary, pt.DataWrapper): # trivial freeze. key_to_frozen_subary[key] = to_tagged_cl_array( subary.data, axes=get_cl_axes_from_pt_axes(subary.axes), tags=subary.tags) - else: - if not isinstance(subary, pt.Array): - raise TypeError(f"{type(self).__name__}.freeze invoked " - f"with non-pytato array of type '{type(array)}'") - + elif isinstance(subary, pt.Array): # Don't be tempted to take shortcuts here, e.g. for empty # arrays, as this will inhibit metadata propagation that # may happen in transform_dag below. See # https://github.com/inducer/arraycontext/pull/167#issuecomment-1151877480 key_to_pt_arrays[key] = subary + else: + raise TypeError( + f"{type(self).__name__}.freeze invoked with an unsupported " + f"array type: got '{type(subary).__name__}', but expected one " + f"of {self.array_types}") # }}} pt_dict_of_named_arrays = pt.make_dict_of_named_arrays( - key_to_pt_arrays) - + key_to_pt_arrays) normalized_expr, bound_arguments = _normalize_pt_expr( pt_dict_of_named_arrays) @@ -352,8 +390,8 @@ def _record_leaf_ary_in_dict(key: Tuple[Any, ...], pt_prg = self._freeze_prg_cache[normalized_expr] except KeyError: try: - transformed_dag, function_name = \ - self._dag_transform_cache[normalized_expr] + transformed_dag, function_name = ( + self._dag_transform_cache[normalized_expr]) except KeyError: transformed_dag = self.transform_dag(normalized_expr) @@ -373,16 +411,16 @@ def _record_leaf_ary_in_dict(key: Tuple[Any, ...], self._dag_transform_cache[normalized_expr] = ( transformed_dag, function_name) + from arraycontext.loopy import _DEFAULT_LOOPY_OPTIONS pt_prg = pt.generate_loopy(transformed_dag, - options=lp.Options(return_dict=True, - no_numpy=True), + options=_DEFAULT_LOOPY_OPTIONS, cl_device=self.queue.device, function_name=function_name) pt_prg = pt_prg.with_transformed_program(self.transform_loopy_program) self._freeze_prg_cache[normalized_expr] = pt_prg else: - transformed_dag, function_name = \ - self._dag_transform_cache[normalized_expr] + transformed_dag, function_name = ( + self._dag_transform_cache[normalized_expr]) assert len(pt_prg.bound_arguments) == 0 evt, out_dict = pt_prg(self.queue, **bound_arguments) @@ -391,47 +429,79 @@ def _record_leaf_ary_in_dict(key: Tuple[Any, ...], key_to_frozen_subary = { **key_to_frozen_subary, - **{k: to_tagged_cl_array(v.with_queue(None), - get_cl_axes_from_pt_axes(transformed_dag[k] - .expr - .axes), - transformed_dag[k].expr.tags) + **{k: to_tagged_cl_array( + v.with_queue(None), + axes=get_cl_axes_from_pt_axes(transformed_dag[k].expr.axes), + tags=transformed_dag[k].expr.tags) for k, v in out_dict.items()} } - def _to_frozen(key: Tuple[Any, ...], ary: ArrayT): + def _to_frozen(key: Tuple[Any, ...], ary) -> TaggableCLArray: key_str = "_ary" + _ary_container_key_stringifier(key) return key_to_frozen_subary[key_str] - return with_array_context(rec_keyed_map_array_container(_to_frozen, - array), - actx=None) + return with_array_context( + rec_keyed_map_array_container(_to_frozen, array), + actx=None) def thaw(self, array): import pytato as pt from .utils import get_pt_axes_from_cl_axes - from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, - to_tagged_cl_array) - import pyopencl.array as cl_array + import arraycontext.impl.pyopencl.taggable_cl_array as tga - def _rec_thaw(ary): - if isinstance(ary, TaggableCLArray): - pass - elif isinstance(ary, cl_array.Array): - ary = to_tagged_cl_array(ary, axes=None, tags=frozenset()) - else: - raise TypeError(f"{type(self).__name__}.thaw expects " - "'TaggableCLArray' or 'cl.array.Array' got " - f"{type(ary)}.") + def _thaw(ary): return pt.make_data_wrapper(ary.with_queue(self.queue), axes=get_pt_axes_from_cl_axes(ary.axes), tags=ary.tags) - return with_array_context(rec_map_array_container(_rec_thaw, array), - actx=self) + return with_array_context( + self._rec_map_container(_thaw, array, (tga.TaggableCLArray,)), + actx=self) + + def tag(self, tags: ToTagSetConvertible, array): + def _tag(ary): + return ary.tagged(_preprocess_array_tags(tags)) + + return self._rec_map_container(_tag, array) + + def tag_axis(self, iaxis, tags: ToTagSetConvertible, array): + def _tag_axis(ary): + return ary.with_tagged_axis(iaxis, tags) + + return self._rec_map_container(_tag_axis, array) # }}} + # {{{ compilation + + def call_loopy(self, program, **kwargs): + import pytato as pt + from pytato.scalar_expr import SCALAR_CLASSES + from pytato.loopy import call_loopy + from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray + + entrypoint = program.default_entrypoint.name + + # {{{ preprocess args + + processed_kwargs = {} + + for kw, arg in sorted(kwargs.items()): + if isinstance(arg, (pt.Array,) + SCALAR_CLASSES): + pass + elif isinstance(arg, TaggableCLArray): + arg = self.thaw(arg) + else: + raise ValueError(f"call_loopy argument '{kw}' expected to be an" + " instance of 'pytato.Array', 'Number' or" + f"'TaggableCLArray', got '{type(arg)}'") + + processed_kwargs[kw] = arg + + # }}} + + return call_loopy(program, processed_kwargs, entrypoint) + def compile(self, f: Callable[..., Any]) -> Callable[..., Any]: from .compile import LazilyPyOpenCLCompilingFunctionCaller return LazilyPyOpenCLCompilingFunctionCaller(self, f) @@ -442,39 +512,30 @@ def transform_dag(self, dag: "pytato.DictOfNamedArrays" dag = pt.transform.materialize_with_mpms(dag) return dag - def tag(self, tags: ToTagSetConvertible, array): - return rec_map_array_container( - lambda x: x.tagged(_preprocess_array_tags(tags)), - array) - - def tag_axis(self, iaxis, tags: ToTagSetConvertible, array): - return rec_map_array_container( - lambda x: x.with_tagged_axis(iaxis, tags), - array) - def einsum(self, spec, *args, arg_names=None, tagged=()): - import pyopencl.array as cla import pytato as pt - from arraycontext.impl.pyopencl.taggable_cl_array import (TaggableCLArray, - to_tagged_cl_array) + import arraycontext.impl.pyopencl.taggable_cl_array as tga + if arg_names is None: arg_names = (None,) * len(args) def preprocess_arg(name, arg): - if isinstance(arg, TaggableCLArray): + if isinstance(arg, tga.TaggableCLArray): ary = self.thaw(arg) - elif isinstance(arg, cla.Array): + elif isinstance(arg, self._frozen_array_types): from warnings import warn - warn("Passing pyopencl.array.Array to einsum will be " - "deprecated in 2023." - " Use `to_tagged_cl_array` to convert the array to" - " TaggableCLArray.", DeprecationWarning, stacklevel=2) - ary = self.thaw(to_tagged_cl_array(arg, - axes=None, - tags=frozenset())) - else: - assert isinstance(arg, pt.Array) + warn(f"Invoking {type(self).__name__}.einsum with" + f" {type(arg).__name__} will be unsupported in 2023. Use" + " `to_tagged_cl_array` to convert instances to TaggableCLArray.", + DeprecationWarning, stacklevel=2) + ary = self.thaw(tga.to_tagged_cl_array(arg)) + elif isinstance(arg, pt.Array): ary = arg + else: + raise TypeError( + f"{type(self).__name__}.einsum invoked with an unsupported " + f"array type: got '{type(arg).__name__}', but expected one " + f"of {self.array_types}") if name is not None: # Tagging Placeholders with naming-related tags is pointless: @@ -493,6 +554,11 @@ def preprocess_arg(name, arg): for name, arg in zip(arg_names, args) ]).tagged(_preprocess_array_tags(tagged)) + def clone(self): + return type(self)(self.queue, self.allocator) + + # }}} + # }}} @@ -521,34 +587,71 @@ def __init__(self, super().__init__(compile_trace_callback=compile_trace_callback) self.array_types = (pt.Array, DeviceArray) - def clone(self): - return type(self)() + @property + def _frozen_array_types(self) -> Tuple[Type, ...]: + from jax.numpy import DeviceArray + return (DeviceArray, ) + + 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 zeros_like(self, ary): + def _zeros_like(array): + return self.zeros(array.shape, array.dtype) - def from_numpy(self, array: Union[np.ndarray, ScalarLike]): + return self._rec_map_container(_zeros_like, ary, default_scalar=0) + + def from_numpy(self, array): import jax import pytato as pt - return pt.make_data_wrapper(jax.device_put(array)) - def to_numpy(self, array): - if np.isscalar(array): - return array + def _from_numpy(ary): + return pt.make_data_wrapper(jax.device_put(ary)) + + return with_array_context( + self._rec_map_container(_from_numpy, array, (np.ndarray,)), + actx=self) + def to_numpy(self, array): import jax - return jax.device_get(self.freeze(array)) - @property - def frozen_array_types(self) -> Tuple[Type, ...]: - from jax.numpy import DeviceArray - return (DeviceArray, ) + def _to_numpy(ary): + return jax.device_get(ary) - def call_loopy(self, program, **kwargs): - raise ValueError(f"{type(self)} does not support calling loopy.") + return with_array_context( + self._rec_map_container(_to_numpy, self.freeze(array)), + actx=None) def freeze(self, array): + if np.isscalar(array): + return array + import pytato as pt from jax.numpy import DeviceArray - from arraycontext.context import ArrayT from arraycontext.container.traversal import rec_keyed_map_array_container from arraycontext.impl.pytato.compile import _ary_container_key_stringifier @@ -557,10 +660,9 @@ def freeze(self, array): key_to_pt_arrays: Dict[str, pt.Array] = {} def _record_leaf_ary_in_dict(key: Tuple[Any, ...], - ary: Union[DeviceArray, pt.Array]): + ary: Union[DeviceArray, pt.Array]) -> None: key_str = "_ary" + _ary_container_key_stringifier(key) array_as_dict[key_str] = ary - return ary rec_keyed_map_array_container(_record_leaf_ary_in_dict, array) @@ -572,12 +674,13 @@ def _record_leaf_ary_in_dict(key: Tuple[Any, ...], elif isinstance(subary, pt.DataWrapper): # trivial freeze. key_to_frozen_subary[key] = subary.data.block_until_ready() - else: - if not isinstance(subary, pt.Array): - raise TypeError(f"{type(self).__name__}.freeze invoked " - f"with non-pytato array of type '{type(array)}'") - + elif isinstance(subary, pt.Array): key_to_pt_arrays[key] = subary + else: + raise TypeError( + f"{type(self).__name__}.freeze invoked with an unsupported " + f"array type: got '{type(subary).__name__}', but expected one " + f"of {self.array_types}") # }}} @@ -593,59 +696,59 @@ def _record_leaf_ary_in_dict(key: Tuple[Any, ...], for k, v in out_dict.items()} } - def _to_frozen(key: Tuple[Any, ...], ary: ArrayT): + def _to_frozen(key: Tuple[Any, ...], ary) -> DeviceArray: key_str = "_ary" + _ary_container_key_stringifier(key) return key_to_frozen_subary[key_str] - return with_array_context(rec_keyed_map_array_container(_to_frozen, - array), - actx=None) + return with_array_context( + rec_keyed_map_array_container(_to_frozen, array), + actx=None) def thaw(self, array): import pytato as pt - from jax.numpy import DeviceArray - def _rec_thaw(ary): - if isinstance(ary, DeviceArray): - pass - else: - raise TypeError(f"{type(self).__name__}.thaw expects " - f"'jax.DeviceArray' got {type(ary)}.") + def _thaw(ary): return pt.make_data_wrapper(ary) - return with_array_context(rec_map_array_container(_rec_thaw, array), - actx=self) + return with_array_context( + self._rec_map_container(_thaw, array, self._frozen_array_types), + actx=self) def compile(self, f: Callable[..., Any]) -> Callable[..., Any]: from .compile import LazilyJAXCompilingFunctionCaller return LazilyJAXCompilingFunctionCaller(self, f) def tag(self, tags: ToTagSetConvertible, array): - import pytato as pt from jax.numpy import DeviceArray - def _rec_tag(ary): + def _tag(ary): if isinstance(ary, DeviceArray): return ary else: - assert isinstance(ary, pt.Array) return ary.tagged(_preprocess_array_tags(tags)) - return rec_map_array_container(_rec_tag, array) + return self._rec_map_container(_tag, array) def tag_axis(self, iaxis, tags: ToTagSetConvertible, array): - import pytato as pt from jax.numpy import DeviceArray - def _rec_tag_axis(ary): + def _tag_axis(ary): if isinstance(ary, DeviceArray): return ary else: - assert isinstance(ary, pt.Array) return ary.with_tagged_axis(iaxis, tags) - return rec_map_array_container(_rec_tag_axis, - array) + return self._rec_map_container(_tag_axis, array) + + # }}} + + # {{{ compilation + + def call_loopy(self, program, **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 pytato as pt @@ -656,9 +759,13 @@ def einsum(self, spec, *args, arg_names=None, tagged=()): def preprocess_arg(name, arg): if isinstance(arg, DeviceArray): ary = self.thaw(arg) - else: - assert isinstance(arg, pt.Array) + elif isinstance(arg, pt.Array): ary = arg + else: + raise TypeError( + f"{type(self).__name__}.einsum invoked with an unsupported " + f"array type: got '{type(arg).__name__}', but expected one " + f"of {self.array_types}") if name is not None: # Tagging Placeholders with naming-related tags is pointless: @@ -677,7 +784,9 @@ def preprocess_arg(name, arg): for name, arg in zip(arg_names, args) ]).tagged(_preprocess_array_tags(tagged)) -# }}} + def clone(self): + return type(self)() +# }}} # vim: foldmethod=marker diff --git a/arraycontext/impl/pytato/compile.py b/arraycontext/impl/pytato/compile.py index 9e92adfd..32823286 100644 --- a/arraycontext/impl/pytato/compile.py +++ b/arraycontext/impl/pytato/compile.py @@ -152,10 +152,10 @@ def _get_arg_id_to_arg_and_arg_id_to_descr(args: Tuple[Any, ...], arg_id_to_descr[arg_id] = ScalarInputDescriptor(np.dtype(type(arg))) elif is_array_container_type(arg.__class__): def id_collector(keys, ary): - arg_id = (kw,) + keys - arg_id_to_arg[arg_id] = ary - arg_id_to_descr[arg_id] = LeafArrayDescriptor(np.dtype(ary.dtype), - ary.shape) + arg_id = (kw,) + keys # noqa: B023 + arg_id_to_arg[arg_id] = ary # noqa: B023 + arg_id_to_descr[arg_id] = LeafArrayDescriptor( # noqa: B023 + np.dtype(ary.dtype), ary.shape) return ary rec_keyed_map_array_container(id_collector, arg) @@ -523,7 +523,7 @@ def _args_to_device_buffers(actx, input_id_to_name_in_program, arg_id_to_arg): elif isinstance(arg, pt.array.DataWrapper): # got a Datawrapper => simply gets its data arg = arg.data - elif isinstance(arg, actx.frozen_array_types): + elif isinstance(arg, actx._frozen_array_types): # got a frozen array => do nothing pass elif isinstance(arg, pt.Array): diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index 0f219d13..d1890f29 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -83,7 +83,8 @@ def full_like(self, ary, fill_value): def _full_like(subary): return pt.full(subary.shape, fill_value, subary.dtype) - return self._new_like(ary, _full_like) + return self._array_context._rec_map_container( + _full_like, ary, default_scalar=fill_value) # }}} diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py index e89663ca..842d108e 100644 --- a/test/test_arraycontext.py +++ b/test/test_arraycontext.py @@ -197,6 +197,84 @@ def _with_actx_dofarray(ary: DOFArray, actx: ArrayContext) -> DOFArray: # type: # }}} +# {{{ nested containers + +@with_container_arithmetic(bcast_obj_array=False, + eq_comparison=False, rel_comparison=False, + _cls_has_array_context_attr=True) +@dataclass_array_container +@dataclass(frozen=True) +class MyContainer: + name: str + mass: Union[DOFArray, np.ndarray] + momentum: np.ndarray + enthalpy: Union[DOFArray, np.ndarray] + + @property + def array_context(self): + if isinstance(self.mass, np.ndarray): + return next(iter(self.mass)).array_context + else: + return self.mass.array_context + + +@with_container_arithmetic( + bcast_obj_array=False, + bcast_container_types=(DOFArray, np.ndarray), + matmul=True, + rel_comparison=True, + _cls_has_array_context_attr=True) +@dataclass_array_container +@dataclass(frozen=True) +class MyContainerDOFBcast: + name: str + mass: Union[DOFArray, np.ndarray] + momentum: np.ndarray + enthalpy: Union[DOFArray, np.ndarray] + + @property + def array_context(self): + if isinstance(self.mass, np.ndarray): + return next(iter(self.mass)).array_context + else: + return self.mass.array_context + + +def _get_test_containers(actx, ambient_dim=2, shapes=50_000): + from numbers import Number + if isinstance(shapes, (Number, tuple)): + shapes = [shapes] + + x = DOFArray(actx, tuple([ + actx.from_numpy(randn(shape, np.float64)) + for shape in shapes])) + + # pylint: disable=unexpected-keyword-arg, no-value-for-parameter + dataclass_of_dofs = MyContainer( + name="container", + mass=x, + momentum=make_obj_array([x] * ambient_dim), + enthalpy=x) + + # pylint: disable=unexpected-keyword-arg, no-value-for-parameter + bcast_dataclass_of_dofs = MyContainerDOFBcast( + name="container", + mass=x, + momentum=make_obj_array([x] * ambient_dim), + enthalpy=x) + + ary_dof = x + ary_of_dofs = make_obj_array([x] * ambient_dim) + mat_of_dofs = np.empty((ambient_dim, ambient_dim), dtype=object) + for i in np.ndindex(mat_of_dofs.shape): + mat_of_dofs[i] = x + + return (ary_dof, ary_of_dofs, mat_of_dofs, dataclass_of_dofs, + bcast_dataclass_of_dofs) + +# }}} + + # {{{ assert_close_to_numpy* def randn(shape, dtype): @@ -341,6 +419,23 @@ def test_array_context_np_like(actx_factory, sym_name, n_args, dtype): assert_close_to_numpy( actx, lambda _np, *_args: getattr(_np, sym_name)(*_args), args) + for c in (42.0,) + _get_test_containers(actx): + result = getattr(actx.np, sym_name)(c) + result = actx.thaw(actx.freeze(result)) + + if sym_name == "zeros_like": + if np.isscalar(result): + assert result == 0.0 + else: + assert actx.to_numpy(actx.np.all(actx.np.equal(result, 0.0))) + elif sym_name == "ones_like": + if np.isscalar(result): + assert result == 1.0 + else: + assert actx.to_numpy(actx.np.all(actx.np.equal(result, 1.0))) + else: + raise ValueError(f"unknown method: '{sym_name}'") + # }}} @@ -671,79 +766,6 @@ def test_array_context_einsum_array_tripleprod(actx_factory, spec): # {{{ array container classes for test -@with_container_arithmetic(bcast_obj_array=False, - eq_comparison=False, rel_comparison=False, - _cls_has_array_context_attr=True) -@dataclass_array_container -@dataclass(frozen=True) -class MyContainer: - name: str - mass: Union[DOFArray, np.ndarray] - momentum: np.ndarray - enthalpy: Union[DOFArray, np.ndarray] - - @property - def array_context(self): - if isinstance(self.mass, np.ndarray): - return next(iter(self.mass)).array_context - else: - return self.mass.array_context - - -@with_container_arithmetic( - bcast_obj_array=False, - bcast_container_types=(DOFArray, np.ndarray), - matmul=True, - rel_comparison=True, - _cls_has_array_context_attr=True) -@dataclass_array_container -@dataclass(frozen=True) -class MyContainerDOFBcast: - name: str - mass: Union[DOFArray, np.ndarray] - momentum: np.ndarray - enthalpy: Union[DOFArray, np.ndarray] - - @property - def array_context(self): - if isinstance(self.mass, np.ndarray): - return next(iter(self.mass)).array_context - else: - return self.mass.array_context - - -def _get_test_containers(actx, ambient_dim=2, shapes=50_000): - from numbers import Number - if isinstance(shapes, (Number, tuple)): - shapes = [shapes] - - x = DOFArray(actx, tuple([ - actx.from_numpy(randn(shape, np.float64)) - for shape in shapes])) - - # pylint: disable=unexpected-keyword-arg, no-value-for-parameter - dataclass_of_dofs = MyContainer( - name="container", - mass=x, - momentum=make_obj_array([x] * ambient_dim), - enthalpy=x) - - # pylint: disable=unexpected-keyword-arg, no-value-for-parameter - bcast_dataclass_of_dofs = MyContainerDOFBcast( - name="container", - mass=x, - momentum=make_obj_array([x] * ambient_dim), - enthalpy=x) - - ary_dof = x - ary_of_dofs = make_obj_array([x] * ambient_dim) - mat_of_dofs = np.empty((ambient_dim, ambient_dim), dtype=object) - for i in np.ndindex(mat_of_dofs.shape): - mat_of_dofs[i] = x - - return (ary_dof, ary_of_dofs, mat_of_dofs, dataclass_of_dofs, - bcast_dataclass_of_dofs) - def test_container_scalar_map(actx_factory): actx = actx_factory()