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: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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 require 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: Iterable[Tag] | Tag | None) → _Self_Taggable

Return a copy of self with the specified tag or tags added to the set of tags. If the resulting set of tags violates the rules on pytools.tag.UniqueTag, an error is raised.

Parameters:

tags – An instance of Tag or an iterable with instances therein.

without_tags(tags: Iterable[Tag] | Tag | None, verify_existence: bool = True) → _Self_Taggable

Return a copy of self without the specified tags.

Parameters:

• tags – An instance of Tag or an iterable with instances therein.

• verify_existence – 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.

Array interface:

__getitem__(slice_spec: int | slice | Array | None | ellipsis | Tuple[int | slice | Array | None | ellipsis, ...]) → Array

Warning

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

T

__mul__()

__rmul__()

__add__()

__radd__()

__sub__()

__rsub__()

__truediv__()

__rtruediv__()

__neg__()

__pos__()

__and__()

__rand__()

__or__()

__ror__()

__xor__()

__rxor__()

__abs__()

conj()

all(axis: int = 0) → Array | number[Any] | int | bool_ | bool | float | complex

Equivalent to pytato.all().

any(axis: int = 0) → Array | number[Any] | int | bool_ | bool | float | complex

Equivalent to pytato.any().

with_tagged_axis(iaxis: int, tags: Sequence[Tag] | Tag) → Array

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

real

imag

Derived attributes:

ndim

class pytato.Axis(tags: FrozenSet[Tag])

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

class pytato.ReductionDescriptor(tags: FrozenSet[Tag])

Records information about a reduction dimension in an Array’.

class pytato.NamedArray(container: AbstractResultWithNamedArrays, name: str, *, axes: AxesT, tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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

__init__(container: AbstractResultWithNamedArrays, name: str, *, axes: AxesT, tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({})) → None

Method generated by attrs for class NamedArray.

class pytato.DictOfNamedArrays(data: Mapping[str, Array], *, tags: FrozenSet[Tag] | None = None)

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

Implements AbstractResultWithNamedArrays.

__init__(data: Mapping[str, Array], *, tags: FrozenSet[Tag] | None = None) → None

Method generated by attrs for class AbstractResultWithNamedArrays.

class pytato.AbstractResultWithNamedArrays(*, tags: FrozenSet[Tag])

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__(*, tags: FrozenSet[Tag]) → None

Method generated by attrs for class AbstractResultWithNamedArrays.

abstract __contains__(name: object) → bool

abstract __getitem__(name: str) → NamedArray

abstract __len__() → 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: Array, x2: Array) → Array

Matrix multiplication.

Parameters:

• x1 – first argument

• x2 – second argument

pytato.roll(a: Array, shift: int, axis: int | None = None) → Array

Roll array elements along a given axis.

Parameters:

• a – input array

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

• axis – axis along which the array is shifted

pytato.transpose(a: Array, axes: Sequence[int] | None = None) → Array

Reverse or permute the axes of an array.

Parameters:

• a – input array

• axes – 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].

pytato.stack(arrays: Sequence[Array], axis: int = 0) → Array

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 – a finite sequence, each of whose elements is an Array of the same shape

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

pytato.concatenate(arrays: Sequence[Array], axis: int = 0) → Array

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 – a finite sequence, each of whose elements is an Array . The arrays are of the same shape except along the axis dimension.

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

pytato.zeros(shape: int | ~numpy.integer[~typing.Any] | ~pytato.array.Array | ~typing.Sequence[int | ~numpy.integer[~typing.Any] | ~pytato.array.Array], dtype: ~typing.Any = <class 'float'>, order: str = 'C') → Array

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

pytato.ones(shape: int | ~numpy.integer[~typing.Any] | ~pytato.array.Array | ~typing.Sequence[int | ~numpy.integer[~typing.Any] | ~pytato.array.Array], dtype: ~typing.Any = <class 'float'>, order: str = 'C') → Array

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

pytato.full(shape: int | integer[Any] | Array | Sequence[int | integer[Any] | Array], fill_value: number[Any] | int | bool_ | bool | float | complex, dtype: Any = None, order: str = 'C') → Array

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

pytato.eye(N: int, M: int | None = None, k: int = 0, dtype: ~typing.Any = <class 'numpy.float64'>) → Array

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

Parameters:

• N – Number of rows in the output matrix

• M – Number of columns in the output matrix. Equal to N if None.

pytato.arange(*args: Any, **kwargs: Any) → Array

arange([start, ]stop, [step, ]dtype=None)

Semantically equivalent to numpy.arange().

pytato.equal(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns (x1 == x2) element-wise.

pytato.not_equal(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns (x1 != x2) element-wise.

pytato.less(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns (x1 < x2) element-wise.

pytato.less_equal(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns (x1 <= x2) element-wise.

pytato.greater(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns (x1 > x2) element-wise.

pytato.greater_equal(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns (x1 >= x2) element-wise.

pytato.logical_or(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

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

pytato.logical_and(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

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

pytato.logical_not(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | bool

Returns the element-wise logical NOT of x.

pytato.where(condition: Array | number[Any] | int | bool_ | bool | float | complex, x: Array | number[Any] | int | bool_ | bool | float | complex | None = None, y: Array | number[Any] | int | bool_ | bool | float | complex | None = None) → Array | number[Any] | int | bool_ | bool | float | complex

Elementwise selector between x and y depending on condition.

pytato.maximum(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

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

pytato.minimum(x1: Array | number[Any] | int | bool_ | bool | float | complex, x2: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

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

pytato.einsum(subscripts: str, *operands: Array, index_to_redn_descr: Mapping[str, ReductionDescriptor] | None = None) → Einsum

Einstein summation subscripts on operands.

pytato.dot(a: Array | number[Any] | int | bool_ | bool | float | complex, b: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

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.

pytato.vdot(a: Array, b: Array) → Array | number[Any] | int | bool_ | bool | float | complex

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.

pytato.broadcast_to(array: Array, shape: Tuple[int | integer[Any] | Array, ...]) → Array

Returns array broadcasted to shape.

pytato.squeeze(array: Array, axis: Collection[int] | None = None) → Array

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

Parameters:

axis – Subset of 1-long axes of array that must be removed. If None all 1-long axes are removed.

pytato.expand_dims(array: Array, axis: Tuple[int, ...] | int) → Array

Reshapes array by adding 1-long axes at axis dimensions of the returned array.

pytato.abs(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.sqrt(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.sin(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.cos(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.tan(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.arcsin(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.arccos(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.arctan(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.conj(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.arctan2(y: Array | number[Any] | int | bool_ | bool | float | complex, x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.sinh(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.cosh(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.tanh(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.exp(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.log(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.log10(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.isnan(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.real(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.imag(x: Array | number[Any] | int | bool_ | bool | float | complex) → Array | number[Any] | int | bool_ | bool | float | complex

pytato.pad.pad(array: Array, pad_width: int | integer[Any] | Sequence[int | integer[Any]], mode: str = 'constant', **kwargs: Any) → Array

Returns an array with padded elements along each axis.

Parameters:

• array – The array to be padded.

• pad_width –

  Number of elements to be padded along each axis. Can be one of:

  • An instance of int denoting the constant number of elements to pad before and after each axis.

  • A tuple of the form (before, after) denoting that before number of padded elements must precede each axis and after number of padded elements must succeed each axis.

  • A sequence with i-th element as the tuple (before_i, after_i) denoting that before_i number of padded elements must precede the i-th axis and after_i number of padded elements must succeed the i-th axis.

• mode –

  An instance of str denoting the values of the padded elements in the returned array. It can be one of:

  • "constant" denoting that the padded elements must be filled with constant entries. See constant_values.

• constant_values –

  Optional argument when operating under "constant" mode. Can be one of:

  • An instance of int denoting the value of every padded element.

  • A tuple of the form (before, after) denoting that every padded element that precedes array’s axes must be set to before and every padded element that succeeds array’s axes must be set to after.

  • A sequence with the i-th element of the form (before_i, after_i) denoting that the padded elements preceding array’s i-th axis must be set to before_i and the padded elements succeeding array’s i-th axis must be set to after_i.

  Defaults to 0.

Note

As of March, 2023 the values of the padded elements that are preceding wrt certain axes and succeeding wrt other axes is undefined as per numpy.pad()‘s spec.

pytato.sum(a: Array, axis: int | Tuple[int, ...] | None = None, initial: Any = 0, axis_to_reduction_descr: Mapping[int, ReductionDescriptor] | None = None) → Array

Sums array a’s elements along the axis axes.

Parameters:

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

• axis – The axes along which the elements are to be sum-reduced. Defaults to all axes of the input array.

• initial – The value returned for an empty array, if supplied. This value also serves as the base value onto which any additional array entries are accumulated.

• axis_to_reduction_descr – A mapping from axis in axis to the corresponding instance of ReductionDescriptor that the IndexLambda is to be instantiated with.

pytato.amin(a: ~pytato.array.Array, axis: int | ~typing.Tuple[int, ...] | None = None, initial: ~typing.Any = <class 'pytato.reductions._NoValue'>, axis_to_reduction_descr: ~typing.Mapping[int, ~pytato.array.ReductionDescriptor] | None = None) → Array

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

Parameters:

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

• axis – The axes along which the elements are to be min-reduced. Defaults to all axes of the input array.

• initial – The value returned for an empty array, if supplied. This value also serves as the base value onto which any additional array entries are accumulated. If not supplied, an ValueError will be raised if the reduction is empty. In that case, the reduction size must not be symbolic.

• axis_to_reduction_descr – A mapping from axis in axis to the corresponding instance of ReductionDescriptor that the IndexLambda is to be instantiated with.

pytato.amax(a: ~pytato.array.Array, axis: int | ~typing.Tuple[int] | None = None, *, initial: ~typing.Any = <class 'pytato.reductions._NoValue'>, axis_to_reduction_descr: ~typing.Mapping[int, ~pytato.array.ReductionDescriptor] | None = None) → Array

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

Parameters:

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

• axis – The axes along which the elements are to be max-reduced. Defaults to all axes of the input array.

• initial – The value returned for an empty array, if supplied. This value also serves as the base value onto which any additional array entries are accumulated. If not supplied, an ValueError will be raised if the reduction is empty. In that case, the reduction size must not be symbolic.

• axis_to_reduction_descr – A mapping from axis in axis to the corresponding instance of ReductionDescriptor that the IndexLambda is to be instantiated with.

pytato.prod(a: Array, axis: int | Tuple[int, ...] | None = None, initial: Any = 1, axis_to_reduction_descr: Mapping[int, ReductionDescriptor] | None = None) → Array

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

Parameters:

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

• axis – The axes along which the elements are to be product-reduced. Defaults to all axes of the input array.

• initial – The value returned for an empty array, if supplied. This value also serves as the base value onto which any additional array entries are accumulated.

• axis_to_reduction_descr – A mapping from axis in axis to the corresponding instance of ReductionDescriptor that the IndexLambda is to be instantiated with.

pytato.all(a: Array, axis: int | Tuple[int, ...] | None = None, axis_to_reduction_descr: Mapping[int, ReductionDescriptor] | None = None) → Array

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

Parameters:

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

• axis – The axes along which the elements are to be product-reduced. Defaults to all axes of the input array.

• axis_to_reduction_descr – A mapping from axis in axis to the corresponding instance of ReductionDescriptor that the IndexLambda is to be instantiated with.

pytato.any(a: Array, axis: int | Tuple[int, ...] | None = None, axis_to_reduction_descr: Mapping[int, ReductionDescriptor] | None = None) → Array

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

Parameters:

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

• axis – The axes along which the elements are to be product-reduced. Defaults to all axes of the input array.

• axis_to_reduction_descr – A mapping from axis in axis to the corresponding instance of ReductionDescriptor that the IndexLambda is to be instantiated with.

class pytato.reductions.ReductionOperation

abstract neutral_element(dtype: dtype[Any]) → Any

abstract __hash__() → int

Return hash(self).

abstract __eq__(other: Any) → bool

Return self==value.

class pytato.reductions.SumReductionOperation

class pytato.reductions.ProductReductionOperation

class pytato.reductions.MaxReductionOperation

class pytato.reductions.MinReductionOperation

class pytato.reductions.AllReductionOperation

class pytato.reductions.AnyReductionOperation

Concrete Array Data

class pytato.array.DataInterface(*args, **kwargs)

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(shape: Tuple[int | integer[Any] | Array, ...], dtype: dtype[Any], expr: Expression, bindings: Mapping[str, Array], var_to_reduction_descr: Mapping[str, ReductionDescriptor], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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.

var_to_reduction_descr

A mapping from reduction variables in expr to their ReductionDescriptor.

with_tagged_reduction(reduction_variable: str, tag: Tag) → IndexLambda

Returns a copy of self with the ReductionDescriptor associated with reduction_variable tagged with tag.

Parameters:

reduction_variable – Name of reduction variable in self that is to be tagged.

class pytato.array.Einsum(access_descriptors: Tuple[Tuple[EinsumAxisDescriptor, ...], ...], args: Tuple[Array, ...], redn_axis_to_redn_descr: Mapping[EinsumReductionAxis, ReductionDescriptor], index_to_access_descr: Mapping[str, EinsumAxisDescriptor], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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.

access_descr_to_index

Mapping from the access descriptors to the index used by the user during the instantiation of the Einsum node. This is a strictly non-semantic attribute and only present to support a friendlier with_tagged_reduction().

with_tagged_reduction(redn_axis: EinsumReductionAxis | str, tag: Tag) → Einsum

Returns a copy of self with the ReductionDescriptor associated with redn_axis tagged with tag.

class pytato.array.Stack(arrays: Tuple[Array, ...], axis: int, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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: Tuple[Array, ...], axis: int, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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: Array, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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: Array, shift: int, axis: int, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

Roll an array along an axis.

shift

Shift amount.

axis

Shift axis.

class pytato.array.AxisPermutation(array: Array, axis_permutation: Tuple[int, ...], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

Permute the axes of an array.

array

axis_permutation

A permutation of the input axes.

class pytato.array.Reshape(array: Array, newshape: Tuple[int | integer[Any] | Array, ...], order: str, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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: Array, indices: Tuple[int | integer[Any] | NormalizedSlice | Array | None | ellipsis, ...], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

Abstract class for all index expressions on an array.

indices

class pytato.array.BasicIndex(array: Array, indices: Tuple[int | integer[Any] | NormalizedSlice | Array | None | ellipsis, ...], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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

class pytato.array.AdvancedIndexInContiguousAxes(array: Array, indices: Tuple[int | integer[Any] | NormalizedSlice | Array | None | ellipsis, ...], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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: Array, indices: Tuple[int | integer[Any] | NormalizedSlice | Array | None | ellipsis, ...], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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(data: DataInterface, shape: Tuple[int | integer[Any] | Array, ...], *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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(shape: Tuple[int | integer[Any] | Array, ...], dtype: dtype[Any], name: str, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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__(shape: Tuple[int | integer[Any] | Array, ...], dtype: dtype[Any], name: str, *, axes: Tuple[Axis, ...], tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({})) → None

Method generated by attrs for class Placeholder.

class pytato.array.SizeParam(name: str, *, tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}), axes: Tuple[Axis, ...] = ())

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: Dict[str, Array], *, tags: FrozenSet[Tag] = frozenset({})) → DictOfNamedArrays

Make a DictOfNamedArrays object.

Parameters:

data – member keys and arrays

pytato.array.make_placeholder(name: str, shape: int | ~numpy.integer[~typing.Any] | ~pytato.array.Array | ~typing.Sequence[int | ~numpy.integer[~typing.Any] | ~pytato.array.Array], dtype: ~typing.Any = <class 'numpy.float64'>, tags: ~typing.FrozenSet[~pytools.tag.Tag] = frozenset({}), axes: ~typing.Tuple[~pytato.array.Axis, ...] | None = None) → Placeholder

Make a Placeholder object.

Parameters:

• name – name of the placeholder array, generated automatically if not given

• shape – shape of the placeholder array

• dtype – dtype of the placeholder array (must be convertible to numpy.dtype, default is numpy.float64)

• tags – implementation tags

pytato.array.make_size_param(name: str, tags: FrozenSet[Tag] = frozenset({})) → SizeParam

Make a SizeParam.

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

Parameters:

• name – name

• tags – implementation tags

pytato.array.make_data_wrapper(data: DataInterface, *, name: str | None = None, shape: int | integer[Any] | Array | Sequence[int | integer[Any] | Array] | None = None, tags: FrozenSet[Tag] = frozenset({}), axes: Tuple[Axis, ...] | None = None) → DataWrapper

Make a DataWrapper.

Parameters:

• data – an instance obeying the DataInterface

• name – an optional name, generated automatically if not given

• shape – optional shape of the array, inferred from data if not given

• tags – implementation tags

Internal API

class pytato.array.EinsumAxisDescriptor

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

class pytato.array.EinsumElementwiseAxis(dim: int)

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

class pytato.array.EinsumReductionAxis(dim: int)

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

class pytato.array.NormalizedSlice(start: ShapeComponent, stop: ShapeComponent, step: IntegralT)

A normalized version of slice. “Normalized” is explained in start and stop.

start

An instance of ShapeComponent. Normalized to satisfy the relation 0 <= start <= (axis_len-1), where axis_len is the length of the axis being sliced.

stop

An instance of ShapeComponent. Normalized to satisfy the relation -1 <= stop <= axis_len, where axis_len is the length of the axis being sliced.

step

Traceback functionality

Please consider these undocumented and subject to change at any time.

pytato.array.set_traceback_tag_enabled(enable: bool = True) → None

Enable or disable the traceback tag.

class pytato.tags._PytatoFrameSummary(filename: str, lineno: int | None, name: str, line: str | None)

Class to store a single call frame, similar to traceback.FrameSummary, but immutable.

class pytato.tags._PytatoStackSummary(frames: Tuple[_PytatoFrameSummary, ...])

Class to store a list of _PytatoFrameSummary call frames, similar to traceback.StackSummary, but immutable.

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

class pytato.array.AxesT

A tuple of Axis objects.

class pytato.array.IntegralT

An integer data type which is a union of integral types of numpy and int.

Functions in Pytato IR

pytato.trace_call(f: ~typing.Callable[[...], ~pytato.function.ReturnT], *args: ~pytato.array.Array, identifier: ~typing.Hashable | None = <class 'pytato.function._Guess'>, **kwargs: ~pytato.array.Array) → ReturnT

Returns the expressions returned after calling f with the arguments args and keyword arguments kwargs. The subexpressions in the returned expressions are outlined (opposite of ‘inlined’) as a FunctionDefinition.

Parameters:

identifier – A hashable object that acts as pytato.tags.FunctionIdentifier.identifier for the FunctionIdentifier tagged to the outlined FunctionDefinition. If None the function definition is not tagged with a FunctionIdentifier tag, if _Guess the function identifier is guessed from f.__name__.

class pytato.function.Call(function: FunctionDefinition, bindings: Mapping[str, Array], *, tags: FrozenSet[Tag])

Records an invocation to a FunctionDefinition.

function

The instance of FunctionDefinition being called by this call site.

bindings

A mapping from the placeholder names of FunctionDefinition to their corresponding parameters in the invocation to function.

class pytato.function.NamedCallResult(call: Call, name: str)

One of the arrays that are returned from a call to FunctionDefinition.

call

The function invocation that led to self.

name

The name by which the returned array is referred to in FunctionDefinition.returns.

class pytato.function.FunctionDefinition(parameters: FrozenSet[str], return_type: ReturnType, returns: Mapping[str, Array], *, tags: FrozenSet[Tag])

A function definition that represents its outputs as instances of Array with the inputs being Placeholders. The outputs of the function can be a single pytato.Array, a tuple of pytato.Arrays or an instance of Dict[str, Array].

parameters

Names of the input Placeholders to the function node. This is a superset of the names of Placeholder instances encountered in returns. Unused parameters are allowed.

return_type

An instance of ReturnType.

returns

The outputs of the function call which are array expressions that depend on the parameters. The keys of the mapping depend on return_type as:

• If the function returns a single pytato.Array, then returns contains a single array expression with "_" as the key.

• If the function returns a tuple of pytato.Arrays, then returns contains entries with the key "_N" mapping the N-th entry of the result-tuple.

• If the function returns a dict mapping identifiers to pytato.Arrays, then returns uses the same mapping.

get_placeholder(name: str) → Placeholder

Returns the instance of pytato.array.Placeholder corresponding to the parameter name in function body.

Note

A FunctionDefinition comes with its own namespace based on parameters. A Mapper-implementer must ensure not to reuse the cached result between the caller’s expressions and a function definition’s expressions to avoid unsound cache hits that could lead to incorrect mappings.

Note

At this point, code generation/execution does not support distributed-memory communication nodes (DistributedSend, DistributedRecv) within function bodies.

class pytato.function.ReturnType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Records the function body’s return type in FunctionDefinition.

class pytato.function.ReturnT

A type variable corresponding to the return type of the function pytato.trace_call().

Raising IndexLambda nodes

class pytato.raising.HighLevelOp

Base class for all high level operations that could be raised from a pytato.array.IndexLambda.

pytato.raising.index_lambda_to_high_level_op(expr: IndexLambda) → HighLevelOp

Returns a HighLevelOp corresponding expr.

Calling loopy kernels in an array expression

class pytato.loopy.LoopyCall(translation_unit: lp.TranslationUnit, bindings: Mapping[str, ArrayOrScalar], entrypoint: str, *, tags: FrozenSet[Tag])

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

class pytato.loopy.LoopyCallResult(name: str, container: LoopyCall, *, axes: AxesT, tags: FrozenSet[Tag], non_equality_tags: FrozenSet[Tag | None] = frozenset({}))

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

pytato.loopy.call_loopy(translation_unit: TranslationUnit, bindings: Dict[str, Array | number[Any] | int | bool_ | bool | float | complex], entrypoint: str | None = None) → LoopyCall

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:

• translation_unit – the translation unit to call.

• bindings – mapping from argument names of translation_unit[entrypoint] to pytato.array.Array.

• entrypoint – the entrypoint of the translation_unit parameter.

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

class pytato.loopy.Tag

See pytools.tag.Tag.

class pytato.loopy.AxesT

See pytato.array.AxesT.

class pytato.loopy.ArrayOrScalar

A Array or a scalar.

class lp.TranslationUnit

See loopy.TranslationUnit.

Attaching metadata to arrays

Pre-Defined Tags

class pytato.tags.ImplementationStrategy

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

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

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: str)

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: str)

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: str)

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

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

class pytato.tags.CreatedAt(traceback: _PytatoStackSummary)

A tag attached to a Array to store the traceback of where it was created.

class pytato.tags.ExpandedDimsReshape(new_dims: Tuple[int, ...])

A tag that can be attached to a Reshape to indicate that the new dimensions created by pytato.expand_dims().

Attr new_dims:

A tuple of the dimensions of the reshaped array that were added.

>>> x = pt.make_placeholder("x", (10, 4), "float64")
>>> pt.expand_dims(x, (0, 2, 4)).tags_of_type(pt.tags.ExpandedDimsReshape)
frozenset({ExpandedDimsReshape(new_dims=(0, 2, 4))})

class pytato.tags.FunctionIdentifier(identifier: Hashable)

A tag that can be attached to a FunctionDefinition node to to describe the function’s identifier. One can use this to refer all instances of FunctionDefinition, for example in transformations.transform.calls.concatenate_calls`.

identifier

class pytato.tags.CallImplementationTag

A tag that can be attached to a Call node to direct a Target how the call site should be lowered.

class pytato.tags.InlineCallTag

A CallImplementationTag that directs the pytato.target.Target to inline the call site.

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: str) → number[Any] | int | bool_ | bool | float | complex | Expression

pytato.scalar_expr.get_dependencies(expression: Any, include_idx_lambda_indices: bool = True) → FrozenSet[str]

Return the set of variable names in an expression.

Parameters:

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

pytato.scalar_expr.substitute(expression: Any, variable_assigments: Mapping[str, Any] | None) → Any

Perform variable substitution in an expression.

Parameters:

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

• variable_assigments – A mapping from variable names to substitutions