Reference

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.Namespace(*args, **kwds)

Represents a mapping from identifier strings to array expressions or None, where None indicates that the name may not be used. (pytato.array.Placeholder instances register their names in this way to avoid ambiguity.)

name_gen
__contains__(name: object) → bool
__getitem__(name: str) → pytato.array.Array
__iter__() → Iterator[str]
__len__() → int
assign(name: str, value: pytato.array.Array) → str

Declare a new array.

Parameters

  • name – a Python identifier

  • value – the array object

Returns

name

copy() → pytato.array.Namespace
ref(name: str) → pytato.array.Array
Returns

An array expression referring to name.

class pytato.Array(tags: Optional[FrozenSet[pytato.array.Tag]] = None)

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

namespace

A (mutable) instance of Namespace containing the names used in the computation. All arrays in a computation share the same namespace.

shape

Identifiers (pymbolic.primitives.Variable) refer to names from namespace. A tuple of integers or pymbolic expressions. Shape may be (at most affinely) symbolic in these identifiers.

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.

tags

A tuple of Tag instances.

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

named(name: str) → pytato.array.Array
tagged(tag: pytato.array.Tag) → pytato.array.Array

Returns a copy of self tagged with tag. If tag is a UniqueTag and other tags of this type are already present, an error is raised.

without_tag(tag: pytato.array.Tag, verify_existence: bool = True) → pytato.array.Array

Array interface:

__getitem__(slice_spec: Union[int, slice, None, ellipsis, Tuple[Union[int, slice, None, ellipsis], …]]) → pytato.array.Array
T
__mul__()
__rmul__()
__add__()
__radd__()
__sub__()
__rsub__()
__truediv__()
__rtruediv__()
__neg__()
__pos__()

Derived attributes:

ndim
class pytato.Tag

Generic metadata, applied to, among other things, instances of Array.

tag_name

A fully qualified DottedName that reflects the class name of the tag.

Instances of this type must be immutable, hashable, picklable, and have a reasonably concise __repr__() of the form dotted.name(attr1=value1, attr2=value2). Positional arguments are not allowed.

__repr__()

Return repr(self).

Note

This mirrors the tagging scheme that loopy is headed towards.

class pytato.UniqueTag

Only one instance of this type of tag may be assigned to a single tagged object.

class pytato.DictOfNamedArrays(*args, **kwds)

A container that maps valid Python identifiers to instances of Array. May occur as a result type of array computations.

__init__(data: Dict[str, pytato.array.Array])

Initialize self. See help(type(self)) for accurate signature.

__contains__(name: object) → bool
__getitem__(name: str) → pytato.array.Array
__iter__() → Iterator[str]
__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: pytato.array.Array, x2: pytato.array.Array) → pytato.array.Array

Matrix multiplication.

Parameters

  • x1 – first argument

  • x2 – second argument

pytato.roll(a: pytato.array.Array, shift: int, axis: Optional[int] = None) → pytato.array.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: pytato.array.Array, axes: Optional[Sequence[int]] = None) → pytato.array.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[pytato.array.Array], axis: int = 0) → pytato.array.Array

Join a sequence of arrays along a new axis.

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

Example:

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

Supporting Functionality

class pytato.DottedName(name_parts: Tuple[str, …])
name_parts

A tuple of strings, each of which is a valid Python identifier. No name part may start with a double underscore.

The name (at least morally) exists in the name space defined by the Python module system. It need not necessarily identify an importable object.

classmethod from_class(argcls: Any) → pytato.array.DottedName

Concrete Array Data

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

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

Pre-Defined Tags

class pytato.array.ImplementAs(strategy: pytato.array.ImplementationStrategy)
strategy
class pytato.array.ImplementationStrategy
class pytato.array.CountNamed(name: str)
name
class pytato.array.ImplStored
class pytato.array.ImplInlined
class pytato.array.ImplDefault

Built-in Expression Nodes

class pytato.array.IndexLambda(namespace: pytato.array.Namespace, expr: pymbolic.primitives.Expression, shape: Tuple[Union[numbers.Number, pymbolic.primitives.Expression], …], dtype: numpy.dtype, bindings: Optional[Dict[str, pytato.array.Array]] = None, tags: Optional[FrozenSet[pytato.array.Tag]] = None)
namespace
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, then in namespace.

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.

is_reference() → bool
class pytato.array.Einsum(tags: Optional[FrozenSet[pytato.array.Tag]] = None)
class pytato.array.MatrixProduct(x1: pytato.array.Array, x2: pytato.array.Array, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

A product of two matrices, or a matrix and a vector.

The semantics of this operation follow PEP 465 [pep465], i.e., the Python matmul (@) operator.

x1
x2
pep465

https://www.python.org/dev/peps/pep-0465/

class pytato.array.LoopyFunction(*args, **kwds)

Note

This should allow both a locally stored kernel and one that’s obtained by importing a dotted name.

class pytato.array.Stack(arrays: Tuple[pytato.array.Array, …], axis: int, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

Join a sequence of arrays along an axis.

arrays

The sequence of arrays to join

axis

The output axis

class pytato.array.AttributeLookup(tags: Optional[FrozenSet[pytato.array.Tag]] = None)

An expression node to extract an array from a DictOfNamedArrays.

Warning

Not yet implemented.

Index Remapping

class pytato.array.IndexRemappingBase(array: pytato.array.Array, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

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: pytato.array.Array, shift: int, axis: int, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

Roll an array along an axis.

shift

Shift amount.

axis

Shift axis.

class pytato.array.AxisPermutation(array: pytato.array.Array, axes: Tuple[int, …], tags: Optional[FrozenSet[pytato.array.Tag]] = None)

Permute the axes of an array.

axes

A permutation of the input axes.

class pytato.array.Reshape(array: pytato.array.Array, tags: Optional[FrozenSet[pytato.array.Tag]] = None)
class pytato.array.Slice(array: pytato.array.Array, begin: Tuple[int, …], size: Tuple[int, …], tags: Optional[FrozenSet[pytato.array.Tag]] = None)

Extracts a slice of constant size from an array.

begin
size

Input Arguments

class pytato.array.InputArgumentBase(namespace: pytato.array.Namespace, name: str, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

Base class for input arguments.

name

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

The name is also implicitly assign()ed in the Namespace.

Note

Creating multiple instances of any input argument with the same name and within the same Namespace is not allowed.

class pytato.array.DataWrapper(namespace: pytato.array.Namespace, name: str, data: pytato.array.DataInterface, shape: Tuple[Union[numbers.Number, pymbolic.primitives.Expression], …], tags: Optional[FrozenSet[pytato.array.Tag]] = None)

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.

class pytato.array.Placeholder(namespace: pytato.array.Namespace, name: str, shape: Tuple[Union[numbers.Number, pymbolic.primitives.Expression], …], dtype: numpy.dtype, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

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

name
__init__(namespace: pytato.array.Namespace, name: str, shape: Tuple[Union[numbers.Number, pymbolic.primitives.Expression], …], dtype: numpy.dtype, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

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

class pytato.array.SizeParam(namespace: pytato.array.Namespace, name: str, tags: Optional[FrozenSet[pytato.array.Tag]] = None)

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

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.ConvertibleToShape
pytato.array.make_dict_of_named_arrays(data: Dict[str, pytato.array.Array]) → pytato.array.DictOfNamedArrays

Make a DictOfNamedArrays object and ensure that all arrays share the same namespace.

Parameters

data – member keys and arrays

pytato.array.make_placeholder(namespace: pytato.array.Namespace, name: str, shape: Union[str, numbers.Number, pymbolic.primitives.Expression, Tuple[Union[int, pymbolic.primitives.Expression, str], …]], dtype: Any, tags: Optional[FrozenSet[pytato.array.Tag]] = None) → pytato.array.Placeholder

Make a Placeholder object.

Parameters

  • namespace – namespace of the placeholder array

  • name – name of the placeholder array

  • shape – shape of the placeholder array

  • dtype – dtype of the placeholder array (must be convertible to numpy.dtype)

  • tags – implementation tags

pytato.array.make_size_param(namespace: pytato.array.Namespace, name: str, tags: Optional[FrozenSet[pytato.array.Tag]] = None) → pytato.array.SizeParam

Make a SizeParam.

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

Parameters

  • namespace – namespace

  • name – name

  • tags – implementation tags

pytato.array.make_data_wrapper(namespace: pytato.array.Namespace, data: pytato.array.DataInterface, name: Optional[str] = None, shape: Optional[Union[str, numbers.Number, pymbolic.primitives.Expression, Tuple[Union[int, pymbolic.primitives.Expression, str], …]]] = None, tags: Optional[FrozenSet[pytato.array.Tag]] = None) → pytato.array.DataWrapper

Make a DataWrapper.

Parameters

  • namespace – namespace

  • 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

Aliases

(This section exists because Sphinx, our documentation tool, can’t (yet) canonicalize type references. Once Sphinx 4.0 is released, we should use the :canonical: option here.)

class pytato.array.Namespace

Should be referenced as pytato.Namespace.

class pytato.array.DottedName

Should be referenced as pytato.DottedName.

class pytato.array.Tag

Should be referenced as pytato.Tag.

class pytato.array.Array

Should be referenced as pytato.Array.

class pytato.array.DictOfNamedArrays

Should be referenced as pytato.DictOfNamedArrays.

Scalar Expressions

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, pymbolic.primitives.Expression].

pytato.scalar_expr.parse(s: str) → Union[numbers.Number, pymbolic.primitives.Expression]
pytato.scalar_expr.get_dependencies(expression: Any) → 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]) → 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

Transforming Computations

class pytato.transform.CopyMapper(namespace: pytato.array.Namespace)
pytato.transform.copy_namespace(source_namespace: pytato.array.Namespace, copy_mapper: pytato.transform.CopyMapper) → pytato.array.Namespace

Copy the elements of namespace into a new namespace.

Parameters

  • source_namespace – The namespace to copy

  • copy_mapper – A mapper that performs copies into a new namespace

Returns

A new namespace containing copies of the items in source_namespace

pytato.transform.copy_dict_of_named_arrays(source_dict: pytato.array.DictOfNamedArrays, copy_mapper: pytato.transform.CopyMapper) → pytato.array.DictOfNamedArrays

Copy the elements of a DictOfNamedArrays into a DictOfNamedArrays with a new namespace.

Parameters

  • source_dict – The DictOfNamedArrays to copy

  • copy_mapper – A mapper that performs copies into a new namespace

Returns

A new DictOfNamedArrays containing copies of the items in source_dict

Generated Executable Programs

class pytato.program.BoundProgram(program: loopy.LoopKernel, bound_arguments: Mapping[str, Any], target: pytato.target.Target)

A wrapper around a loopy kernel for execution.

program

The underlying loopy.LoopKernel.

target

The code generation target.

bound_arguments

A map from names to pre-bound kernel arguments.

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

Call self as a function.

class pytato.program.BoundPyOpenCLProgram(program: loopy.LoopKernel, bound_arguments: Mapping[str, Any], target: pytato.target.Target, queue: Optional[pyopencl.CommandQueue])

A wrapper around a loopy kernel for execution with pyopencl.

queue

A pyopencl command queue.

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

Convenience function for launching a pyopencl computation.

Code Generation Targets

class pytato.target.Target

An abstract code generation target.

get_loopy_target() → loopy.TargetBase

Return the corresponding loopy target.

bind_program(program: loopy.LoopKernel, bound_arguments: Mapping[str, Any]) → pytato.program.BoundProgram

Create a pytato.program.BoundProgram for this code generation target.

Parameters

  • program – the loopy kernel

  • bound_arguments – a mapping from argument names to outputs

class pytato.target.PyOpenCLTarget(queue: Optional[pyopencl.CommandQueue] = None)

A pyopencl code generation target.

queue

The pyopencl command queue, or None.

References

class pytato.codegen.DictOfNamedArrays

Should be referenced as pytato.DictOfNamedArrays.

class pytato.codegen.DataInterface

Should be referenced as pytato.array.DataInterface.

Generating Code

pytato.generate_loopy(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays], target: Optional[pytato.target.Target] = None, options: Optional[loopy.options.Options] = None) → pytato.program.BoundProgram

Code generation entry point.

Parameters

  • result – Outputs of the computation.

  • target – Code generation target.

  • options – Code generation options for the kernel.

Returns

A wrapped generated loopy kernel

Code Generation Internals

class pytato.codegen.CodeGenPreprocessor(namespace: pytato.array.Namespace)

A mapper that preprocesses graphs to simplify code generation.

The following node simplifications are performed:

Source Node Type

Target Node Type

DataWrapper

Placeholder

Roll

IndexLambda

AxisPermutation

IndexLambda

Slice

IndexLambda
class pytato.codegen.LoopyExpressionContext(state: pytato.codegen.CodeGenState, num_indices: int, _depends_on: FrozenSet[str] = <factory>, local_namespace: Mapping[str, pytato.array.Array] = <factory>, reduction_bounds: Dict[str, Tuple[Union[numbers.Number, pymbolic.primitives.Expression], Union[numbers.Number, pymbolic.primitives.Expression]]] = <factory>)

Mutable state used while generating loopy expressions. Wraps CodeGenState with more expression-specific information.

This data is passed through InlinedExpressionGenMapper via arguments, and is also used by ImplementedResult.to_loopy_expression() to retrieve contextual data.

state

The CodeGenState.

local_namespace

A (read-only) local name mapping used for name lookup when generating code.

num_indices

The number of indices of the form _0, _1, allowed in the expression.

depends_on

The set of statement IDs that need to be included in loopy.InstructionBase.depends_on.

reduction_bounds

A mapping from inames to reduction bounds in the expression.

update_depends_on(other: FrozenSet[str]) → None
lookup(name: str) → pytato.array.Array
class pytato.codegen.ImplementedResult

Generated code for a node in the computation graph (i.e., an array expression).

to_loopy_expression(indices: Tuple[Union[numbers.Number, pymbolic.primitives.Expression], …], expr_context: pytato.codegen.LoopyExpressionContext) → Union[numbers.Number, pymbolic.primitives.Expression]

Return a loopy expression for this result.

Parameters

  • indices – symbolic expressions for the indices of the array

  • expr_context –

    the associated expression context. The fields are treated as follows:

    • depends_on is populated with any dependencies needed for the generated expression.

    • reduction_bounds is populated with reduction bounds for the reduction inames in the returned expression. If reduction_bounds is nonempty, then the returned inames are ensured to be disjoint from those present.

class pytato.codegen.StoredResult(name: str, num_indices: int, depends_on: FrozenSet[str])

An array expression generated as a loopy array.

See also: pytato.array.ImplStored.

class pytato.codegen.InlinedResult(expr: Union[numbers.Number, pymbolic.primitives.Expression], num_indices: int, reduction_bounds: Dict[str, Tuple[Union[numbers.Number, pymbolic.primitives.Expression], Union[numbers.Number, pymbolic.primitives.Expression]]], depends_on: FrozenSet[str])

An array expression generated as a loopy expression containing inlined sub-expressions.

See also: pytato.array.ImplInlined.

class pytato.codegen.SubstitutionRuleResult
class pytato.codegen.CodeGenState(namespace: Mapping[str, pytato.array.Array], _kernel: loopy.kernel.LoopKernel, results: Dict[pytato.array.Array, pytato.codegen.ImplementedResult])

A container for data kept by CodeGenMapper.

namespace

The (global) namespace

kernel

The partial loopy.LoopKernel being built.

results

A mapping from pytato.Array instances to instances of ImplementedResult.

var_name_gen
insn_id_gen
update_kernel(kernel: loopy.kernel.LoopKernel) → None
class pytato.codegen.CodeGenMapper

A mapper for generating code for nodes in the computation graph.

class pytato.codegen.InlinedExpressionGenMapper(codegen_mapper: pytato.codegen.CodeGenMapper)

A mapper for generating loopy expressions with inlined sub-expressions.

The inputs to this mapper are scalar expression as found in pytato.array.IndexLambda, or expressions that are compatible (e.g., shape expressions).

The outputs of this mapper are scalar expressions suitable for wrapping in InlinedResult.

pytato.codegen.domain_for_shape(dim_names: Tuple[str, …], shape: Tuple[Union[numbers.Number, pymbolic.primitives.Expression], …], reductions: Dict[str, Tuple[Union[numbers.Number, pymbolic.primitives.Expression], Union[numbers.Number, pymbolic.primitives.Expression]]]) → islpy._isl.BasicSet

Create an islpy.BasicSet that expresses an appropriate index domain for an array of (potentially symbolic) shape shape having reduction dimensions reductions.

Parameters

  • dim_names – A tuple of strings, the names of the axes. These become set dimensions in the returned domain.

  • shape – A tuple of constant or quasi-affine pymbolic expressions. The variables in these expressions become parameter dimensions in the returned set. Must have the same length as dim_names.

  • reductions – A map from reduction inames to (lower, upper) bounds (as half-open integer ranges). The variables in the bounds become parameter dimensions in the returned set.

pytato.codegen.get_loopy_temporary(name: str, expr: pytato.array.Array) → loopy.kernel.data.TemporaryVariable
pytato.codegen.add_store(name: str, expr: pytato.array.Array, result: pytato.codegen.ImplementedResult, state: pytato.codegen.CodeGenState, output_to_temporary: bool = False) → str

Add an instruction that stores to a variable in the kernel.

Parameters

  • name – name of the output array, which is created

  • expr – the Array to store

  • result – the corresponding ImplementedResult

  • state – code generation state

  • output_to_temporary – whether to generate an output argument (default) or a temporary variable

Returns

the id of the generated instruction

pytato.codegen.rename_reductions(loopy_expr: Union[numbers.Number, pymbolic.primitives.Expression], loopy_expr_context: pytato.codegen.LoopyExpressionContext, var_name_gen: Callable[[str], str]) → Union[numbers.Number, pymbolic.primitives.Expression]

Rename the reduction variables in loopy_expr and loopy_expr_context using the callable var_name_gen.

pytato.codegen.normalize_outputs(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → pytato.array.DictOfNamedArrays

Convert outputs of a computation to the canonical form.

Performs a conversion to DictOfNamedArrays if necessary.

Parameters

result – Outputs of the computation.

pytato.codegen.get_initial_codegen_state(namespace: pytato.array.Namespace, target: pytato.target.Target, options: Optional[loopy.options.Options]) → pytato.codegen.CodeGenState
class pytato.codegen.PreprocessResult(outputs: DictOfNamedArrays, bound_arguments: Dict)
pytato.codegen.preprocess(outputs: pytato.array.DictOfNamedArrays) → pytato.codegen.PreprocessResult

Preprocess a computation for code generation.

Graph Visualization

pytato.get_dot_graph(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → str

Return a string in the dot language depicting the graph of the computation of result.

Parameters

result – Outputs of the computation (cf. pytato.generate_loopy()).

pytato.show_dot_graph(result: Union[pytato.array.Array, pytato.array.DictOfNamedArrays]) → None

Show a graph representing the computation of result in a browser.

Parameters

result – Outputs of the computation (cf. pytato.generate_loopy()).