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() or get_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 a pyopencl.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 and DOFArray will not allow them.)

• For the lazily-evaluating array context based on pytato, “thawing” corresponds to the creation of a symbolic “handle” (specifically, a pytato.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 interface of an array context

class arraycontext.ArrayContext

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 empty(shape: int | Tuple[int, ...], dtype: dtype[Any]) → Array

abstract zeros(shape: int | Tuple[int, ...], dtype: dtype[Any]) → Array

empty_like(ary: Array) → Array

zeros_like(ary: Array) → Array

abstract from_numpy(array: ndarray | ArrayContainer | int | float | complex | generic) → Array | ArrayContainer | int | float | complex | generic

Returns:

the numpy.ndarray array converted to the array context’s array type. The returned array will be thaw()ed. When working with array containers each leaf must be an ndarray 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

Returns:

an numpy.ndarray for each array recognized by the context. The input array must be 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 ndarray leaving the container structure intact.

abstract call_loopy(program: loopy.TranslationUnit, **kwargs: Any) → Dict[str, Array]

Execute the loopy program program on the arguments kwargs.

program is a loopy.LoopKernel or loopy.TranslationUnit. It is expected to not yet be transformed for execution speed. It must have loopy.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

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 from numpy.

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.ArrayContainers.

array_types

A tuple of types that are the valid array classes the context can operate on. However, it is not necessary that all the ArrayContext‘s operations would be legal for the types in array_types.

abstract freeze(array: ArrayOrContainerOrScalarT) → ArrayOrContainerOrScalarT

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 to thaw() it in a different one, as long as that context understands the array format.

abstract thaw(array: ArrayOrContainerOrScalarT) → ArrayOrContainerOrScalarT

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

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() and thaw().

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

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

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]

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 as actx.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.

arraycontext.tag_axes(actx: ArrayContext, dim_to_tags: Mapping[int, Iterable[Tag] | Tag | None], ary: ArrayT) → ArrayT

Return a copy of ary with the axes in dim_to_tags tagged with their corresponding tags. Equivalent to repeated application of ArrayContext.tag_axis().

Types and Type Variables for Arrays and Containers

class arraycontext.Array(*args, **kwargs)

A Protocol for the array type supported by ArrayContext.

This is meant to aid in typing annotations. For a explicit list of supported types see ArrayContext.array_types.

shape

size

dtype

__getitem__

class arraycontext.ArrayT

A type variable with a lower bound of Array.

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

Implementations of the Array Context Abstraction

Array context based on pyopencl.array

class arraycontext.PyOpenCLArrayContext(queue: pyopencl.CommandQueue, allocator: pyopencl.tools.AllocatorBase | None = None, wait_event_queue_length: int | None = None, force_device_scalars: bool = False)

A ArrayContext that uses pyopencl.array.Array instances for its base array class.

context

A pyopencl.Context.

queue

A pyopencl.CommandQueue.

allocator

A PyOpenCL memory allocator. Can also be None (default) or False to use the default allocator. Please note that running with the default allocator allocates and deallocates OpenCL buffers directly. If lots of arrays are created (e.g. as results of computation), the associated cost may become significant. Using e.g. pyopencl.tools.MemoryPool as the allocator can help avoid this cost.

transform_loopy_program(t_unit)

class arraycontext.impl.pyopencl.taggable_cl_array.TaggableCLArray(cq, shape, dtype, order='C', allocator=None, data=None, offset=0, strides=None, events=None, _flags=None, _fast=False, _size=None, _context=None, _queue=None, axes=None, tags=frozenset({}))

A pyopencl.array.Array with additional metadata. This is used by PytatoPyOpenCLArrayContext to preserve tags for data while frozen, and also in a similar capacity by PyOpenCLArrayContext.

axes

A tuple of instances of Axis, with one Axis for each dimension of the array.

tags

A frozenset of pytools.tag.Tag. Typically intended to record application-specific metadata to drive the optimizations in arraycontext.PyOpenCLArrayContext.transform_loopy_program().

class arraycontext.impl.pyopencl.taggable_cl_array.Axis(tags: FrozenSet[Tag])

Records the tags corresponding to a dimensions of TaggableCLArray.

arraycontext.impl.pyopencl.taggable_cl_array.to_tagged_cl_array(ary: Array, axes: Tuple[Axis, ...] | None = None, tags: FrozenSet[Tag] = frozenset({})) → TaggableCLArray

Returns a TaggableCLArray that is constructed from the data in ary along with the metadata from axes and tags. If ary is already a TaggableCLArray, the new tags and axes are added to the existing ones.

Parameters:

axes – An instance of Axis for each dimension of the array. If passed None, then initialized to a pytato.Axis with no tags attached for each dimension.

Lazy/Deferred evaluation array context based on pytato

A pytato-based array context defers the evaluation of an array until its frozen. The execution contexts for the evaluations are specific to an ArrayContext type. For ex. PytatoPyOpenCLArrayContext uses pyopencl to JIT-compile and execute the array expressions.

Following pytato-based array context are provided:

class arraycontext.PytatoPyOpenCLArrayContext(queue: CommandQueue, allocator=None, *, use_memory_pool: bool | None = None, compile_trace_callback: Callable[[Any, str, Any], None] | None = None, _force_svm_arg_limit: int | None = None)

A ArrayContext that uses pytato data types to represent the arrays targeting OpenCL for offloading operations.

queue

A pyopencl.CommandQueue.

allocator

A pyopencl memory allocator. Can also be None (default) or False to use the default allocator.

__init__(queue: CommandQueue, allocator=None, *, use_memory_pool: bool | None = None, compile_trace_callback: Callable[[Any, str, Any], None] | None = None, _force_svm_arg_limit: int | None = None) → None

Parameters:

compile_trace_callback – A function of three arguments (what, stage, ir), where what identifies the object being compiled, stage is a string describing the compilation pass, and ir is an object containing the intermediate representation. This interface should be considered unstable.

transform_dag(dag: pytato.DictOfNamedArrays) → pytato.DictOfNamedArrays

Returns a transformed version of dag. Sub-classes are supposed to override this method to implement context-specific transformations on dag (most likely to perform domain-specific optimizations). Every pytato DAG that is compiled to a GPU-kernel is passed through this routine.

Parameters:

dag – An instance of pytato.DictOfNamedArrays

Returns:

A transformed version of dag.

compile(f: Callable[[...], Any]) → Callable[[...], Any]

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 as actx.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.

class arraycontext.PytatoJAXArrayContext(*, compile_trace_callback: Callable[[Any, str, Any], None] | None = None)

An arraycontext that uses pytato to represent the thawed state of the arrays and compiles the expressions using pytato.target.python.JAXPythonTarget.

Compiling a Python callable (Internal)

class arraycontext.impl.pytato.compile.BaseLazilyCompilingFunctionCaller(actx: ~arraycontext.impl.pytato._BasePytatoArrayContext, f: ~typing.Callable[[...], ~typing.Any], program_cache: ~typing.Dict[~typing.Mapping[~typing.Tuple[~typing.Any, ...], ~arraycontext.impl.pytato.compile.AbstractInputDescriptor], ~arraycontext.impl.pytato.compile.CompiledFunction] = <factory>)

Records a side-effect-free callable f that can be specialized for the input types with which __call__() is invoked.

f

The callable that will be called to obtain pytato DAGs.

__call__(*args: Any, **kwargs: Any) → Any

Returns the result of f’s function application on args.

Before applying f, it is compiled to a pytato DAG that would apply f with args in a lazy-sense. The intermediary pytato DAG for args is memoized in self.

class arraycontext.impl.pytato.compile.LazilyPyOpenCLCompilingFunctionCaller(actx: ~arraycontext.impl.pytato._BasePytatoArrayContext, f: ~typing.Callable[[...], ~typing.Any], program_cache: ~typing.Dict[~typing.Mapping[~typing.Tuple[~typing.Any, ...], ~arraycontext.impl.pytato.compile.AbstractInputDescriptor], ~arraycontext.impl.pytato.compile.CompiledFunction] = <factory>)

class arraycontext.impl.pytato.compile.LazilyJAXCompilingFunctionCaller(actx: ~arraycontext.impl.pytato._BasePytatoArrayContext, f: ~typing.Callable[[...], ~typing.Any], program_cache: ~typing.Dict[~typing.Mapping[~typing.Tuple[~typing.Any, ...], ~arraycontext.impl.pytato.compile.AbstractInputDescriptor], ~arraycontext.impl.pytato.compile.CompiledFunction] = <factory>)

class arraycontext.impl.pytato.compile.CompiledFunction

A callable which captures the pytato.target.BoundProgram resulting from calling f with a given set of input types, and generating loopy IR from it.

pytato_program

input_id_to_name_in_program

A mapping from input id to the placeholder name in CompiledFunction.pytato_program. Input id is represented as the position of f’s argument augmented with the leaf array’s key if the argument is an array container.

abstract __call__(arg_id_to_arg) → Any

Parameters:

arg_id_to_arg – Mapping from input id to the passed argument. See CompiledFunction.input_id_to_name_in_program for input id’s representation.

class arraycontext.impl.pytato.compile.FromArrayContextCompile

Tagged to the entrypoint kernel of every translation unit that is generated by compile().

Typically this tag serves as a branch condition in implementing a specialized transform strategy for kernels compiled by compile().

Array context based on jax.numpy

class arraycontext.EagerJAXArrayContext

A ArrayContext that uses jax.Array instances for its base array class and performs all array operations eagerly. See PytatoJAXArrayContext for a lazier version.

Note

JAX stores a global configuration state in jax.config. Callers are expected to maintain those. Most important for scientific computing workloads being jax_enable_x64.

numpy coverage

This is a list of functionality implemented by arraycontext.ArrayContext.np.

Note

Only functions and methods that have at least one implementation are listed.

Array creation routines

Function	PyOpenCLArrayContext	EagerJAXArrayContext	PytatoPyOpenCLArrayContext	PytatoJAXArrayContext
numpy.empty_like()	Yes	Yes	Yes	Yes
numpy.ones_like()	Yes	Yes	Yes	Yes
numpy.zeros_like()	Yes	Yes	Yes	Yes
numpy.full_like()	Yes	Yes	Yes	Yes
numpy.copy()	Yes	Yes	No	No

Array manipulation routines

Function	PyOpenCLArrayContext	EagerJAXArrayContext	PytatoPyOpenCLArrayContext	PytatoJAXArrayContext
numpy.reshape()	Yes	Yes	Yes	Yes
numpy.ravel()	Yes	Yes	Yes	Yes
numpy.transpose()	No	Yes	Yes	Yes
numpy.broadcast_to()	No	Yes	Yes	Yes
numpy.concatenate()	Yes	Yes	Yes	Yes
numpy.stack()	Yes	Yes	Yes	Yes

Linear algebra

Function	PyOpenCLArrayContext	EagerJAXArrayContext	PytatoPyOpenCLArrayContext	PytatoJAXArrayContext
numpy.vdot()	Yes	Yes	No	No

Logic Functions

Function	PyOpenCLArrayContext	EagerJAXArrayContext	PytatoPyOpenCLArrayContext	PytatoJAXArrayContext
numpy.all()	Yes	Yes	Yes	Yes
numpy.any()	Yes	Yes	Yes	Yes
numpy.greater	Yes	Yes	Yes	Yes
numpy.greater_equal	Yes	Yes	Yes	Yes
numpy.less	Yes	Yes	Yes	Yes
numpy.less_equal	Yes	Yes	Yes	Yes
numpy.equal	Yes	Yes	Yes	Yes
numpy.not_equal	Yes	Yes	Yes	Yes

Mathematical functions

Function	PyOpenCLArrayContext	EagerJAXArrayContext	PytatoPyOpenCLArrayContext	PytatoJAXArrayContext
numpy.sin	Yes	Yes	Yes	Yes
numpy.cos	Yes	Yes	Yes	Yes
numpy.tan	Yes	Yes	Yes	Yes
numpy.arcsin	Yes	Yes	Yes	Yes
numpy.arccos	Yes	Yes	Yes	Yes
numpy.arctan	Yes	Yes	Yes	Yes
numpy.arctan2	Yes	Yes	Yes	Yes
numpy.sinh	Yes	Yes	Yes	Yes
numpy.cosh	Yes	Yes	Yes	Yes
numpy.tanh	Yes	Yes	Yes	Yes
numpy.floor	Yes	Yes	Yes	Yes
numpy.ceil	Yes	Yes	Yes	Yes
numpy.sum()	Yes	Yes	Yes	Yes
numpy.exp	Yes	Yes	Yes	Yes
numpy.log	Yes	Yes	Yes	Yes
numpy.log10	Yes	Yes	Yes	Yes
numpy.real()	Yes	Yes	Yes	Yes
numpy.imag()	Yes	Yes	Yes	Yes
numpy.conjugate	Yes	Yes	Yes	Yes
numpy.maximum	Yes	Yes	Yes	Yes
numpy.amax()	Yes	Yes	Yes	Yes
numpy.minimum	Yes	Yes	Yes	Yes
numpy.amin()	Yes	Yes	Yes	Yes
numpy.sqrt	Yes	Yes	Yes	Yes
numpy.absolute	Yes	Yes	Yes	Yes
numpy.fabs	Yes	Yes	Yes	Yes