The Array Context Abstraction¶
An array context is an abstraction that helps you dispatch between multiple
implementations of numpy
-like \(n\)-dimensional arrays.
Freezing and thawing¶
One of the central concepts introduced by the array context formalism is
the notion of freeze()
and
thaw()
. Each array handled by the array context
is either “thawed” or “frozen”. Unlike the real-world concept of freezing and
thawing, these operations leave the original array alone; instead, a semantically
separate array in the desired state is returned.
“Thawed” arrays are associated with an array context. They use that context to carry out operations (arithmetic, function calls).
“Frozen” arrays are static data. They are not associated with an array context, and no operations can be performed on them.
Freezing and thawing may be used to move arrays from one array context to another,
as long as both array contexts use identical in-memory data representation.
Otherwise, a common format must be agreed upon, for example using
numpy
through to_numpy()
and
from_numpy()
.
Usage guidelines¶
Here are some rules of thumb to use when dealing with thawing and freezing:
Any array that is stored for a long time needs to be frozen. “Memoized” data (cf.
pytools.memoize()
and friends) is a good example of long-lived data that should be frozen.Within a function, if the user did not supply an array context, then any data returned to the user should be frozen.
Note that array contexts need not necessarily be passed as a separate argument. Passing thawed data as an argument to a function suffices to supply an array context. The array context can be extracted from a thawed argument using, e.g.,
get_container_context_opt()
orget_container_context_recursively()
.
What does this mean concretely?¶
Freezing and thawing are abstract names for concrete operations. It may be helpful to understand what these operations mean in the concrete case of various actual array contexts:
Each
PyOpenCLArrayContext
is associated with apyopencl.CommandQueue
. In order to operate on array data, such a command queue is necessary; it is the main means of synchronization between the host program and the compute device. “Thawing” here means associating an array with a command queue, and “freezing” means ensuring that the array data is fully computed in memory and decoupling the array from the command queue. It is not valid to “mix” arrays associated with multiple queues within an operation: if it were allowed, a dependent operation might begin computing before an input is fully available. (Since bugs of this nature would be very difficult to find,pyopencl.array.Array
andDOFArray
will not allow them.)For the lazily-evaluating array context based on
pytato
, “thawing” corresponds to the creation of a symbolic “handle” (specifically, apytato.array.DataWrapper
) representing the array that can then be used in computation, and “freezing” corresponds to triggering (code generation and) evaluation of an array expression that has been built up by the user (using, e.g.pytato.generate_loopy()
).
The ArrayContext
Interface¶
- class arraycontext.ArrayContext[source]¶
- Canonical:
arraycontext.ArrayContext
An interface that allows software implementing a numerical algorithm (such as
Discretization
) to create and interact with arrays without knowing their types.Added in version 2020.2.
- abstract from_numpy(array: ndarray | ArrayContainer | int | float | complex | generic) Array | ArrayContainer | int | float | complex | generic [source]¶
- Returns:
the
numpy.ndarray
array converted to the array context’s array type. The returned array will bethaw()
ed. When working with array containers each leaf must be anndarray
or scalar, which is then converted to the context’s array type leaving the container structure intact.
- abstract to_numpy(array: Array | ArrayContainer | int | float | complex | generic) ndarray | ArrayContainer | int | float | complex | generic [source]¶
- Returns:
an
numpy.ndarray
for each array recognized by the context. The input array must bethaw()
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 anndarray
leaving the container structure intact.
- abstract call_loopy(t_unit: loopy.TranslationUnit, **kwargs: Any) Dict[str, Array] [source]¶
Execute the
loopy
program program on the arguments kwargs.program is a
loopy.LoopKernel
orloopy.TranslationUnit
. It is expected to not yet be transformed for execution speed. It must haveloopy.Options.return_dict
set.- Returns:
a
dict
of outputs from the program, each an array understood by the context.
- einsum(spec: str, *args: Array, arg_names: Tuple[str, ...] | None = None, tagged: Iterable[Tag] | Tag | None = ()) Array [source]¶
Computes the result of Einstein summation following the convention in
numpy.einsum()
.- Parameters:
spec – a string denoting the subscripts for summation as a comma-separated list of subscript labels. This follows the usual
numpy.einsum()
convention. Note that the explicit indicator -> for the precise output form is required.args – a sequence of array-like operands, whose order matches the subscript labels provided by spec.
arg_names – an optional iterable of string types denoting the names of the args. If None, default names will be generated.
tagged – an optional sequence of
pytools.tag.Tag
objects specifying the tags to be applied to the operation.
- Returns:
the output of the einsum
loopy
program
- np¶
Provides access to a namespace that serves as a work-alike to
numpy
. The actual level of functionality provided is up to the individual array context implementation, however the functions and objects available under this namespace must not behave differently fromnumpy
.As a baseline, special functions available through
loopy
(e.g.sin
,exp
) are accessible through this interface. A full list of implemented functionality is given in numpy coverage.Callables accessible through this namespace vectorize over object arrays, including
arraycontext.ArrayContainer
s.
- array_types¶
A
tuple
of types that are the valid array classes the context can operate on. However, it is not necessary that all theArrayContext
‘s operations are legal for the types in array_types. Note that this tuple is only intended for use withisinstance()
. Other uses are not allowed. This allows for ‘types’ with overriddenclass.__instancecheck__()
.
- abstract freeze(array: ArrayOrContainerOrScalarT) ArrayOrContainerOrScalarT [source]¶
Return a version of the context-defined array array that is ‘frozen’, i.e. suitable for long-term storage and reuse. Frozen arrays do not support arithmetic. For example, in the context of
Array
, this might mean stripping the array of an associated command queue, whereas in a lazily-evaluated context, it might mean that the array is evaluated and stored.Freezing makes the array independent of this
ArrayContext
; it is permitted tothaw()
it in a different one, as long as that context understands the array format.
- abstract thaw(array: ArrayOrContainerOrScalarT) ArrayOrContainerOrScalarT [source]¶
Take a ‘frozen’ array and return a new array representing the data in array that is able to perform arithmetic and other operations, using the execution resources of this context. In the context of
Array
, this might mean that the array is equipped with a command queue, whereas in a lazily-evaluated context, it might mean that the returned array is a symbol bound to the data in array.The returned array may not be used with other contexts while thawed.
- freeze_thaw(array: ArrayOrContainerOrScalarT) ArrayOrContainerOrScalarT [source]¶
Evaluate an input array or container to “frozen” data return a new “thawed” array or container representing the evaluation result that is ready for use. This is a shortcut for calling
freeze()
andthaw()
.This method can be useful in array contexts backed by, e.g.
pytato
, to force the evaluation of a built-up array expression (and thereby avoid reevaluations for expressions built on the array).
- abstract tag(tags: Iterable[Tag] | Tag | None, array: ArrayOrContainerOrScalarT) ArrayOrContainerOrScalarT [source]¶
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. When working with array containers, the tags are applied to each leaf of the container.
See Metadata (“tags”) for Arrays and Array Axes as well as application-specific metadata types.
Added in version 2021.2.
- abstract tag_axis(iaxis: int, tags: Iterable[Tag] | Tag | None, array: ArrayOrContainerOrScalarT) ArrayOrContainerOrScalarT [source]¶
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. When working with array containers, the tags are applied to each leaf of the container.
See Metadata (“tags”) for Arrays and Array Axes as well as application-specific metadata types.
Added in version 2021.2.
- compile(f: Callable[[...], Any]) Callable[[...], Any] [source]¶
Compiles f for repeated use on this array context. f is expected to be a pure function performing an array computation.
Control flow statements (
if
,while
) that might take different paths depending on the data lead to undefined behavior and are illegal. Any data-dependent control flow must be expressed via array functions, such asactx.np.where
.f may be called on placeholder data, to obtain a representation of the computation performed, or it may be called as part of the actual computation, on actual data. If f is called on placeholder data, it may be called only once (or a few times).
- Parameters:
f – the function executing the computation.
- Returns:
a function with the same signature as f.
Types and Type Variables for Arrays and Containers¶
- class arraycontext.Array(*args, **kwargs)[source]¶
A
Protocol
for the array type supported byArrayContext
.This is meant to aid in typing annotations. For a explicit list of supported types see
ArrayContext.array_types
.- shape¶
- size¶
- dtype¶
- class arraycontext.ScalarLike¶
A type annotation for scalar types commonly usable with arrays.
See also ArrayContainer
and ArrayOrContainerT
.
- class arraycontext.ArrayOrContainer¶
- class arraycontext.ArrayOrContainerT¶
A type variable with a lower bound of
ArrayOrContainer
.
- class arraycontext.ArrayOrContainerOrScalar¶
- class arraycontext.ArrayOrContainerOrScalarT¶
A type variable with a lower bound of
ArrayOrContainerOrScalar
.
Internal typing helpers (do not import)¶
This is only here because the documentation tool wants it.
- class arraycontext.context.SelfType¶
Canonical locations for type annotations¶
- class arraycontext.context.ArrayT¶
- Canonical:
arraycontext.ArrayT
- class arraycontext.context.ArrayOrContainerT¶
- Canonical:
arraycontext.ArrayOrContainerT
- class arraycontext.context.ArrayOrContainerOrScalarT¶
- Canonical:
arraycontext.ArrayOrContainerOrScalarT