PytatoArrayContext.freeze: support container freezes

arraycontext/__init__.py

@@ -61,7 +61,7 @@
  thaw, freeze,
  flatten, unflatten, flat_size_and_dtype,
  from_numpy, to_numpy,
- outer)
+ outer, with_array_context)
 
 from .impl.pyopencl import PyOpenCLArrayContext
 from .impl.pytato import (PytatoPyOpenCLArrayContext,
@@ -101,7 +101,7 @@
  "rec_map_reduce_array_container", "rec_multimap_reduce_array_container",
  "thaw", "freeze",
  "flatten", "unflatten", "flat_size_and_dtype",
- "from_numpy", "to_numpy",
+ "from_numpy", "to_numpy", "with_array_context",
  "outer",
 
  "PyOpenCLArrayContext", "PytatoPyOpenCLArrayContext",

arraycontext/container/traversal.py

@@ -65,13 +65,15 @@
 
 from typing import Any, Callable, Iterable, List, Optional, Union, Tuple
 from functools import update_wrapper, partial, singledispatch
+from warnings import warn
 
 import numpy as np
 
 from arraycontext.context import ArrayContext, Array, _ScalarLike
 from arraycontext.container import (
  ArrayT, ContainerT, ArrayOrContainerT, NotAnArrayContainerError,
- serialize_container, deserialize_container)
+ serialize_container, deserialize_container,
+ get_container_context_recursively_opt)
 
 
 # {{{ array container traversal helpers
@@ -519,7 +521,6 @@ def _reduce_wrapper(ary: ContainerT, iterable: Iterable[Tuple[Any, Any]]) -> Any
 
 # {{{ freeze/thaw
 
-@singledispatch
 def freeze(
  ary: ArrayOrContainerT,
  actx: Optional[ArrayContext] = None) -> ArrayOrContainerT:
@@ -533,23 +534,33 @@ def freeze(
 
  See :meth:`ArrayContext.thaw`.
  """
- try:
- iterable = serialize_container(ary)
- except NotAnArrayContainerError:
- if actx is None:
- raise TypeError(
- f"cannot freeze arrays of type {type(ary).__name__} "
- "when actx is not supplied. Try calling actx.freeze "
- "directly or supplying an array context")
- else:
- return actx.freeze(ary)
+
+ if actx is None:
+ warn("Calling freeze(ary) without specifying actx is deprecated, explicitly"
+ " call actx.freeze(ary) instead. This will stop working in 2023.",
+ DeprecationWarning, stacklevel=2)
+
+ actx = get_container_context_recursively_opt(ary)
  else:
- return deserialize_container(ary, [
- (key, freeze(subary, actx=actx)) for key, subary in iterable
- ])
+ warn("Calling freeze(ary, actx) is deprecated, call actx.freeze(ary)"
+ " instead. This will stop working in 2023.",
+ DeprecationWarning, stacklevel=2)
+
+ if __debug__:
+ rec_actx = get_container_context_recursively_opt(ary)
+ if (rec_actx is not None) and (rec_actx is not actx):
+ raise ValueError("Supplied array context does not agree with"
+ " the one obtained by traversing 'ary'.")
+
+ if actx is None:
+ raise TypeError(
+ f"cannot freeze arrays of type {type(ary).__name__} "
+ "when actx is not supplied. Try calling actx.freeze "
+ "directly or supplying an array context")
+
+ return actx.freeze(ary)
 
 
-@singledispatch
 def thaw(ary: ArrayOrContainerT, actx: ArrayContext) -> ArrayOrContainerT:
  r"""Thaws recursively by going through all components of the
  :class:`ArrayContainer` *ary*.
@@ -570,14 +581,41 @@ def thaw(ary: ArrayOrContainerT, actx: ArrayContext) -> ArrayOrContainerT:
  in :mod:`meshmode`. This was necessary because
  :func:`~functools.singledispatch` only dispatches on the first argument.
  """
+ warn("Calling thaw(ary, actx) is deprecated, call actx.thaw(ary) instead."
+ " This will stop working in 2023.",
+ DeprecationWarning, stacklevel=2)
+
+ if __debug__:
+ rec_actx = get_container_context_recursively_opt(ary)
+ if rec_actx is not None:
+ raise ValueError("cannot thaw a container that already has an array"
+ " context.")
+
+ return actx.thaw(ary)
+
+# }}}
+
+
+# {{{ with_array_context
+
+@singledispatch
+def with_array_context(ary: ArrayOrContainerT,
+ actx: Optional[ArrayContext]) -> ArrayOrContainerT:
+ """
+ Recursively associates *actx* to all the components of *ary*.
+
+ Array container types may use :func:`functools.singledispatch` ``.register``
+ to register container-specific implementations. See `this issue
+ <https://github.com/inducer/arraycontext/issues/162>`__ for discussion of
+ the future of this functionality.
+ """
  try:
  iterable = serialize_container(ary)
  except NotAnArrayContainerError:
- return actx.thaw(ary)
+ return ary
  else:
- return deserialize_container(ary, [
- (key, thaw(subary, actx)) for key, subary in iterable
- ])
+ return deserialize_container(ary, [(key, with_array_context(subary, actx))
+ for key, subary in iterable])
 
 # }}}
 
@@ -842,7 +880,10 @@ def _to_numpy_with_check(subary: Any) -> Any:
  f"array of type '{type(subary).__name__}' not in "
  f"supported types {actx.array_types}")
 
- return rec_map_array_container(_to_numpy_with_check, ary)
+ return rec_map_array_container(_to_numpy_with_check,
+ # do a freeze first, if 'actx' supports
+ # container-wide freezes
+ thaw(freeze(ary, actx), actx))
 
 # }}}
 

arraycontext/impl/jax/__init__.py

@@ -32,6 +32,8 @@
 from typing import Union, Callable, Any
 from pytools.tag import ToTagSetConvertible
 from arraycontext.context import ArrayContext, _ScalarLike
+from arraycontext.container.traversal import (with_array_context,
+ rec_map_array_container)
 
 
 class EagerJAXArrayContext(ArrayContext):
@@ -84,10 +86,21 @@ def call_loopy(self, t_unit, **kwargs):
  " operations using ArrayContext.np.")
 
  def freeze(self, array):
- return array.block_until_ready()
+ 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)}.")
+ return ary.block_until_ready()
+
+ return with_array_context(rec_map_array_container(_rec_freeze, array),
+ actx=None)
 
  def thaw(self, array):
- return array
+ return with_array_context(array, actx=self)
 
  # }}}
 

arraycontext/impl/pyopencl/__init__.py

@@ -36,7 +36,8 @@
 from pytools.tag import ToTagSetConvertible
 
 from arraycontext.context import ArrayContext, _ScalarLike
-from arraycontext.container.traversal import rec_map_array_container
+from arraycontext.container.traversal import (rec_map_array_container,
+ with_array_context)
 
 
 if TYPE_CHECKING:
@@ -207,27 +208,41 @@ def call_loopy(self, t_unit, **kwargs):
  for name, ary in result.items()}
 
  def freeze(self, array):
- array.finish()
- return array.with_queue(None)
+ import pyopencl.array as cl_array
+
+ 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
 
- if isinstance(array, TaggableCLArray):
- return array.with_queue(self.queue)
- elif isinstance(array, 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(array, axes=None, tags=frozenset())
- .with_queue(self.queue))
- else:
- raise ValueError("array should be a cl.array.Array,"
- f" got '{type(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)
 
  # }}}