Array Expressions

Note

Expression trees based on this package are picklable as long as no non-picklable data (e.g. pyopencl.array.Array) is referenced from DataWrapper.

Array Interface

class pytato.Array(axes, tags)[source]

A base class (abstract interface + supplemental functionality) for lazily evaluating array expressions. The interface seeks to maximize numpy compatibility, though not at all costs.

Objects of this type are hashable and support structural equality comparison (and are therefore immutable).

Note

Hashability and equality testing does break numpy compatibility, purposefully so.

FIXME: Point out our equivalent for numpy’s ==.

shape

A tuple of integers or scalar-shaped Arrays. Array-valued shape components may be (at most affinely) symbolic in terms of SizeParams.

Note

Affine-ness is mainly required by code generation for IndexLambda, but IndexLambda is used to produce references to named arrays. Since any array that needs to be referenced in this way needs to obey this restriction anyway, a decision was made to requir the same of all array expressions.

dtype

An instance of numpy.dtype.

axes

A tuple of Axis instances. One corresponding to each dimension of the array.

tags

A frozenset of pytools.tag.Tag instances.

Motivation: RDF triples (subject: implicitly the array being tagged, predicate: the tag, object: the arg).

Inherits from pytools.tag.Taggable.

tagged(tags)[source]

Return a copy of self with the specified tag or tags unioned. If tags is a pytools.tag.UniqueTag and other tags of this type are already present, an error is raised Assumes self.copy(tags=<NEW VALUE>) is implemented.

Parameters

tags (Union[Iterable[Tag], Tag, None]) – An instance of Tag or an iterable with instances therein.

Return type

~T_co

without_tags(tags, verify_existence=True)[source]

Return a copy of self without the specified tags. self.copy(tags=<NEW VALUE>) is implemented.

Parameters
  • tags (Union[Iterable[Tag], Tag, None]) – An instance of Tag or an iterable with instances therein.

  • verify_existence (bool) – If set to True, this method raises an exception if not all tags specified for removal are present in the original set of tags. Default True.

Return type

~T_co

Array interface:

__getitem__(slice_spec)[source]

Warning

Out-of-bounds accesses via Array indices are undefined behavior and may pass silently.

Return type

Array

T
__mul__()[source]
__rmul__()[source]
__add__()[source]
__radd__()[source]
__sub__()[source]
__rsub__()[source]
__truediv__()[source]
__rtruediv__()[source]
__neg__()[source]
__pos__()[source]
__and__()[source]
__rand__()[source]
__or__()[source]
__ror__()[source]
__xor__()[source]
__rxor__()[source]
__abs__()[source]
conj()[source]
all(axis=0)[source]

Equivalent to pytato.all().

Return type

Union[Array, Number, int, bool_, bool]

any(axis=0)[source]

Equivalent to pytato.any().

Return type

Union[Array, Number, int, bool_, bool]

with_tagged_axis(iaxis, tags)[source]

Returns a copy of self with iaxis-th axis tagged with tags.

Return type

Array

real
Return type

Union[Array, Number, int, bool_, bool]

imag
Return type

Union[Array, Number, int, bool_, bool]

Derived attributes:

ndim
class pytato.Axis(tags)[source]

A type for recording the information about an Array’s axis.

class pytato.NamedArray(container, name, axes, tags=frozenset({}))[source]

An entry in a AbstractResultWithNamedArrays. Holds a reference back to thecontaining instance as well as the name by which self is known there.

__init__(container, name, axes, tags=frozenset({}))[source]

Constructor for all objects that possess a tags attribute.

Parameters

tags (FrozenSet[Tag]) – a frozenset of Tag objects. Tags can be modified via the tagged() and without_tags() routines. Input checking of tags should be performed before creating a Taggable instance, using check_tag_uniqueness().

class pytato.DictOfNamedArrays(data)[source]

A container of named results, each of which can be computed as an array expression provided to the constructor.

Implements AbstractResultWithNamedArrays.

__init__(data)[source]
class pytato.AbstractResultWithNamedArrays[source]

An abstract array computation that results in multiple Arrays, each named. The way in which the values of these arrays are computed is determined by concrete subclasses of this class, e.g. pytato.loopy.LoopyCall or DictOfNamedArrays.

__init__()
abstract __contains__(name)[source]
Return type

bool

abstract __getitem__(name)[source]
Return type

NamedArray

abstract __len__()[source]
Return type

int

Note

This container deliberately does not implement arithmetic.

NumPy-Like Interface

These functions generally follow the interface of the corresponding functions in numpy, but not all NumPy features may be supported.

pytato.matmul(x1, x2)[source]

Matrix multiplication.

Parameters
  • x1 (Array) – first argument

  • x2 (Array) – second argument

Return type

Array

pytato.roll(a, shift, axis=None)[source]

Roll array elements along a given axis.

Parameters
  • a (Array) – input array

  • shift (int) – the number of places by which elements are shifted

  • axis (Optional[int]) – axis along which the array is shifted

Return type

Array

pytato.transpose(a, axes=None)[source]

Reverse or permute the axes of an array.

Parameters
  • a (Array) – input array

  • axes (Optional[Sequence[int]]) – if specified, a permutation of [0, 1, ..., a.ndim-1]. Defaults to range(a.ndim)[::-1]. The returned axis at index i corresponds to the input axis axes[i].

Return type

Array

pytato.stack(arrays, axis=0)[source]

Join a sequence of arrays along a new axis.

The axis parameter specifies the position of the new axis in the result.

Example:

>>> import pytato as pt
>>> arrays = [pt.zeros(3)] * 4
>>> pt.stack(arrays, axis=0).shape
(4, 3)
Parameters
  • arrays (Sequence[Array]) – a finite sequence, each of whose elements is an Array of the same shape

  • axis (int) – the position of the new axis, which will have length len(arrays)

Return type

Array

pytato.concatenate(arrays, axis=0)[source]

Join a sequence of arrays along an existing axis.

Example:

>>> import pytato as pt
>>> arrays = [pt.zeros(3)] * 4
>>> pt.concatenate(arrays, axis=0).shape
(12,)
Parameters
  • arrays (Sequence[Array]) – a finite sequence, each of whose elements is an Array . The arrays are of the same shape except along the axis dimension.

  • axis (int) – The axis along which the arrays will be concatenated.

Return type

Array

pytato.zeros(shape, dtype=<class 'float'>, order='C')[source]

Returns an array of shape shape with all entries equal to 0.

Return type

Array

pytato.ones(shape, dtype=<class 'float'>, order='C')[source]

Returns an array of shape shape with all entries equal to 1.

Return type

Array

pytato.full(shape, fill_value, dtype, order='C')[source]

Returns an array of shape shape with all entries equal to fill_value.

Return type

Array

pytato.eye(N, M=None, k=0, dtype=<class 'numpy.float64'>)[source]

Returns a 2D-array with ones on the k-th diagonal

Parameters
  • N (int) – Number of rows in the output matrix

  • M (Optional[int]) – Number of columns in the output matrix. Equal to N if None.

Return type

Array

pytato.equal(x1, x2)[source]

Returns (x1 == x2) element-wise.

Return type

Union[Array, bool]

pytato.not_equal(x1, x2)[source]

Returns (x1 != x2) element-wise.

Return type

Union[Array, bool]

pytato.less(x1, x2)[source]

Returns (x1 < x2) element-wise.

Return type

Union[Array, bool]

pytato.less_equal(x1, x2)[source]

Returns (x1 <= x2) element-wise.

Return type

Union[Array, bool]

pytato.greater(x1, x2)[source]

Returns (x1 > x2) element-wise.

Return type

Union[Array, bool]

pytato.greater_equal(x1, x2)[source]

Returns (x1 >= x2) element-wise.

Return type

Union[Array, bool]

pytato.logical_or(x1, x2)[source]

Returns the element-wise logical OR of x1 and x2.

Return type

Union[Array, bool]

pytato.logical_and(x1, x2)[source]

Returns the element-wise logical AND of x1 and x2.

Return type

Union[Array, bool]

pytato.logical_not(x)[source]

Returns the element-wise logical NOT of x.

Return type

Union[Array, bool]

pytato.where(condition, x=None, y=None)[source]

Elementwise selector between x and y depending on condition.

Return type

Union[Array, Number, int, bool_, bool]

pytato.maximum(x1, x2)[source]

Returns the elementwise maximum of x1, x2. x1, x2 being array-like objects that could be broadcasted together. NaNs are propagated.

Return type

Union[Array, Number, int, bool_, bool]

pytato.minimum(x1, x2)[source]

Returns the elementwise minimum of x1, x2. x1, x2 being array-like objects that could be broadcasted together. NaNs are propagated.

Return type

Union[Array, Number, int, bool_, bool]

pytato.einsum(subscripts, *operands)[source]

Einstein summation subscripts on operands.

Return type

Einsum

pytato.dot(a, b)[source]

For 1-dimensional arrays a and b computes their inner product. See numpy.dot() for behavior in the case when a and b aren’t single-dimensional arrays.

Return type

Union[Array, Number, int, bool_, bool]

pytato.vdot(a, b)[source]

Returns the dot-product of conjugate of a with b. If the input arguments are multi-dimensional arrays, they are ravel-ed first and then their vdot is computed.

Return type

Union[Array, Number, int, bool_, bool]

pytato.broadcast_to(array, shape)[source]

Returns array broadcasted to shape.

Return type

Array

pytato.squeeze(array)[source]

Remove single-dimensional entries from the shape of an array.

Return type

Array

pytato.abs(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.sqrt(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.sin(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.cos(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.tan(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.arcsin(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.arccos(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.arctan(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.conj(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.arctan2(y, x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.sinh(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.cosh(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.tanh(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.exp(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.log(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.log10(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.isnan(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.real(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.imag(x)[source]
Return type

Union[Array, Number, int, bool_, bool]

pytato.sum(a, axis=None)[source]

Sums array a’s elements along the axis axes.

Parameters
  • a (Array) – The pytato.Array on which to perform the reduction.

  • axis (Union[int, Tuple[int], None]) – The axes along which the elements are to be sum-reduced. Defaults to all axes of the input array.

Return type

Array

pytato.amin(a, axis=None)[source]

Returns the min of array a’s elements along the axis axes.

Parameters
  • a (Array) – The pytato.Array on which to perform the reduction.

  • axis (Union[int, Tuple[int], None]) – The axes along which the elements are to be min-reduced. Defaults to all axes of the input array.

Return type

Array

pytato.amax(a, axis=None)[source]

Returns the max of array a’s elements along the axis axes.

Parameters
  • a (Array) – The pytato.Array on which to perform the reduction.

  • axis (Union[int, Tuple[int], None]) – The axes along which the elements are to be max-reduced. Defaults to all axes of the input array.

Return type

Array

pytato.prod(a, axis=None)[source]

Returns the product of array a’s elements along the axis axes.

Parameters
  • a (Array) – The pytato.Array on which to perform the reduction.

  • axis (Union[int, Tuple[int], None]) – The axes along which the elements are to be product-reduced. Defaults to all axes of the input array.

Return type

Array

pytato.all(a, axis=None)[source]

Returns the logical-and array a’s elements along the axis axes.

Parameters
  • a (Array) – The pytato.Array on which to perform the reduction.

  • axis (Union[int, Tuple[int], None]) – The axes along which the elements are to be product-reduced. Defaults to all axes of the input array.

Return type

Array

pytato.any(a, axis=None)[source]

Returns the logical-or of array a’s elements along the axis axes.

Parameters
  • a (Array) – The pytato.Array on which to perform the reduction.

  • axis (Union[int, Tuple[int], None]) – The axes along which the elements are to be product-reduced. Defaults to all axes of the input array.

Return type

Array

Concrete Array Data

class pytato.array.DataInterface(*args, **kwargs)[source]

A protocol specifying the minimal interface requirements for concrete array data supported by DataWrapper.

See typing.Protocol for more information about protocols.

Code generation targets may impose additional restrictions on the kinds of concrete array data they support.

shape
dtype

Built-in Expression Nodes

class pytato.array.IndexLambda(expr, shape, dtype, bindings, axes, tags=frozenset({}))[source]

Represents an array that can be computed by evaluating expr for every value of the input indices. The input indices are represented by Variables with names _1, _2, and so on.

expr

A scalar-valued pymbolic expression such as a[_1] + b[_2, _1].

Identifiers in the expression are resolved, in order, by lookups in bindings.

Scalar functions in this expression must be identified by a dotted name representing a Python object (e.g. pytato.c99.sin).

bindings

A dict mapping strings that are valid Python identifiers to objects implementing the Array interface, making array expressions available for use in expr.

class pytato.array.Einsum(access_descriptors, args, axes, tags=frozenset({}))[source]

An array expression using the Einstein summation convention. See numpy.einsum() for a similar construct.

Note

Use pytato.einsum() to create this type of expression node in user code.

access_descriptors

A tuple of access_descriptor for each arg in Einsum.args. An access_descriptor is a tuple of EinsumAxisDescriptor denoting how each axis of the argument will be operated in the einstein summation.

args

A tuple of array over which the Einstein summation is being performed.

class pytato.array.Stack(arrays, axis, axes, tags=frozenset({}))[source]

Join a sequence of arrays along a new axis.

arrays

The sequence of arrays to join

axis

The output axis

class pytato.array.Concatenate(arrays, axis, axes, tags=frozenset({}))[source]

Join a sequence of arrays along an existing axis.

arrays

An instance of tuple of the arrays to join. The arrays must have same shape except for the dimension corresponding to axis.

axis

The axis along which the arrays are to be concatenated.

Index Remapping

class pytato.array.IndexRemappingBase(array, axes, tags=frozenset({}))[source]

Base class for operations that remap the indices of an array.

Note that index remappings can also be expressed via IndexLambda.

array

The input Array

class pytato.array.Roll(array, shift, axis, axes, tags=frozenset({}))[source]

Roll an array along an axis.

shift

Shift amount.

axis

Shift axis.

class pytato.array.AxisPermutation(array, axis_permutation, axes, tags=frozenset({}))[source]

Permute the axes of an array.

array
axis_permutation

A permutation of the input axes.

class pytato.array.Reshape(array, newshape, order, axes, tags=frozenset({}))[source]

Reshape an array.

array

The array to be reshaped

newshape

The output shape

order

Output layout order, either C or F.

class pytato.array.IndexBase(array, indices, axes, tags=frozenset({}))[source]

Abstract class for all index expressions on an array.

indices
class pytato.array.BasicIndex(array, indices, axes, tags=frozenset({}))[source]

An indexing expression with all indices being either an int or slice.

class pytato.array.AdvancedIndexInContiguousAxes(array, indices, axes, tags=frozenset({}))[source]

An indexing expression with at least one Array index and all the advanced indices (i.e. scalars/array) appearing contiguously in IndexBase.indices.

The reason for the existence of this class and AdvancedIndexInNoncontiguousAxes is that numpy treats those two cases differently, and we’re bound to follow its precedent.

class pytato.array.AdvancedIndexInNoncontiguousAxes(array, indices, axes, tags=frozenset({}))[source]

An indexing expression with advanced indices (i.e. scalars/arrays) appearing non-contiguously in IndexBase.indices.

The reason for the existence of this class and AdvancedIndexInContiguousAxes is that numpy treats those two cases differently, and we’re bound to follow its precedent.

Input Arguments

class pytato.array.InputArgumentBase(axes, tags)[source]

Base class for input arguments.

Note

Creating multiple instances of any input argument with the same name in an expression is not allowed.

class pytato.array.DataWrapper(name, data, shape, axes, tags=frozenset({}))[source]

Takes concrete array data and packages it to be compatible with the Array interface.

data

A concrete array (containing data), given as, for example, a numpy.ndarray, or a pyopencl.array.Array. This must offer shape and dtype attributes but is otherwise considered opaque. At evaluation time, its type must be understood by the appropriate execution backend.

Starting with the construction of the DataWrapper, this array may not be updated in-place.

shape

The shape of the array is represented separately from array to allow symbolic shapes to be used, and to ease pytato’s job in recognizing shapes of arrays as equal. For example, if the shape of data is (3, 4), and shape is (nrows, ncolumns), then this represents a (global) constraint that that nrows == 3 and ncolumns == 4. Arithmetic and other operations in pytato do not currently resolve these constraints to assess whether shapes match, and thus it is important that a canonical (symbolic) form of the shape tuple is used.

name

An (optional, string) name by which this object can be identified. Hypothetically, this could be used to ‘swap out’ the data captured here, but that functionality is not currently available.

Note

Since we cannot compare instances of DataInterface being wrapped, a DataWrapper instances compare equal to themselves (i.e. the very same instance).

class pytato.array.Placeholder(name, shape, dtype, axes, tags=frozenset({}))[source]

A named placeholder for an array whose concrete value is supplied by the user during evaluation.

name

The name by which a value is supplied for the argument once computation begins.

__init__(name, shape, dtype, axes, tags=frozenset({}))[source]

Should not be called directly. Use make_placeholder() instead.

class pytato.array.SizeParam(name, axes=(), tags=frozenset({}))[source]

A named placeholder for a scalar that may be used as a variable in symbolic expressions for array sizes.

name

The name by which a value is supplied for the argument once computation begins.

User-Facing Node Creation

Node constructors such as Placeholder.__init__ and __init__ offer limited input validation (in favor of faster execution). Node creation from outside pytato should use the following interfaces:

class pytato.array.ShapeComponent
class pytato.array.ConvertibleToShape
pytato.array.make_dict_of_named_arrays(data)[source]

Make a DictOfNamedArrays object.

Parameters

data (Dict[str, Array]) – member keys and arrays

Return type

DictOfNamedArrays

pytato.array.make_placeholder(name, shape, dtype, tags=frozenset({}), axes=None)[source]

Make a Placeholder object.

Parameters
Return type

Placeholder

pytato.array.make_size_param(name, tags=frozenset({}))[source]

Make a SizeParam.

Size parameters may be used as variables in symbolic expressions for array sizes.

Parameters
Return type

SizeParam

pytato.array.make_data_wrapper(data, name=None, shape=None, tags=frozenset({}), axes=None)[source]

Make a DataWrapper.

Parameters
Return type

DataWrapper

Internal API

class pytato.array.EinsumAxisDescriptor[source]

Records the access pattern of iterating over an array’s axis in a Einsum.

class pytato.array.ElementwiseAxis(dim)[source]

Describes an elementwise access pattern of an array’s axis. In terms of the nomenclature used by IndexLambda, ElementwiseAxis(dim=1) would correspond to indexing the array’s axis as _1 in the expression.

class pytato.array.ReductionAxis(dim)[source]

Describes a reduction access pattern of an array’s axis. In terms of the nomenclature used by IndexLambda, ReductionAxis(dim=0) would correspond to indexing the array’s axis as _r0 in the expression.

Internal stuff that is only here because the documentation tool wants it

class pytato.array.T_co

A covariant type variable used in, e.g. pytools.tag.Taggable.copy().

Calling loopy kernels in an array expression

class pytato.loopy.LoopyCall(translation_unit, bindings, entrypoint)[source]

An array expression node representing a call to an entrypoint in a loopy translation unit.

class pytato.loopy.LoopyCallResult(loopy_call, name, axes, tags=frozenset({}))[source]

Named array for LoopyCall’s result. Inherits from NamedArray.

pytato.loopy.call_loopy(translation_unit, bindings, entrypoint=None)[source]

Invokes an entry point of a loopy.TranslationUnit on the array inputs as specified by bindings.

Restrictions on the structure of translation_unit[entrypoint]:

  • array arguments of translation_unit[entrypoint] must either be either input-only or output-only.

  • all input-only arguments of translation_unit[entrypoint] must appear in bindings.

  • all output-only arguments of translation_unit[entrypoint] must appear in bindings.

  • if translation_unit has been declared with multiple entrypoints, entrypoint can not be None.

Parameters
Return type

LoopyCall

Attaching metadata to arrays

Pre-Defined Tags

class pytato.tags.ImplementationStrategy[source]

Metadata to be attached to pytato.Array to convey information to a pytato.target.Target on how it is supposed to be lowered.

class pytato.tags.ImplStored[source]

An ImplementationStrategy that is tagged to an Array to indicate that the Target must allocate a buffer for storing all the array’s elements, and, all the users of the array must read from that buffer.

class pytato.tags.ImplInlined[source]

An ImplementationStrategy that is tagged to an Array to indicate that the Target should inline the tagged array’s expression into its users.

class pytato.tags.CountNamed(name)[source]

Tagged to a bool-dtyped Array A. If A appears as one of the indices in IndexBase, the number of True values in A is assigned to a variable named name in the generated code.

name
class pytato.tags.Named(name)[source]

Tagged to an Array to indicate the Target that if the tagged array is allocated to a variable, then it must be named name.

name
class pytato.tags.PrefixNamed(prefix)[source]

Tagged to an Array to indicate the Target that if the tagged array is allocated to a variable, then its name must begin with prefix.

prefix
class pytato.tags.AssumeNonNegative[source]

A tag attached to a Array to indicate the Target that all entries of the tagged array are non-negative.

Scalar Expressions

Scalar expressions occur in shapes and in the pytato.array.IndexLambda.

pytato.scalar_expr.ScalarExpression

A type for scalar-valued symbolic expressions. Expressions are composable and manipulable via pymbolic.

Concretely, this is an alias for Union[Number, np.bool_, bool, pymbolic.primitives.Expression].

pytato.scalar_expr.parse(s)[source]
Return type

Union[Number, int, bool_, bool, Expression]

pytato.scalar_expr.get_dependencies(expression, include_idx_lambda_indices=True)[source]

Return the set of variable names in an expression.

Parameters

expression (Any) – A scalar expression, or an expression derived from such (e.g., a tuple of scalar expressions)

Return type

FrozenSet[str]

pytato.scalar_expr.substitute(expression, variable_assigments)[source]

Perform variable substitution in an expression.

Parameters
  • expression (Any) – A scalar expression, or an expression derived from such (e.g., a tuple of scalar expressions)

  • variable_assigments (Optional[Mapping[str, Any]]) – A mapping from variable names to substitutions

Return type

Any