From b89cf2d59cd3f09a4d04ea8125eb23abcd26ca6a Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Mon, 29 Nov 2021 19:01:53 -0600 Subject: [PATCH 01/10] Rename EinsumAxisDescriptor subclasses. In order to avoid coming off as something that is more generally applicable to any pt.Array node. --- pytato/analysis/__init__.py | 7 ++++--- pytato/array.py | 30 +++++++++++++++--------------- pytato/codegen.py | 6 +++--- pytato/target/python/numpy_like.py | 4 ++-- 4 files changed, 24 insertions(+), 23 deletions(-) diff --git a/pytato/analysis/__init__.py b/pytato/analysis/__init__.py index e9c80d578..d6cd67d53 100644 --- a/pytato/analysis/__init__.py +++ b/pytato/analysis/__init__.py @@ -227,7 +227,8 @@ def is_einsum_similar_to_subscript(expr: Einsum, subscripts: str) -> bool: would compute the same result as *expr*. """ - from pytato.array import ElementwiseAxis, ReductionAxis, EinsumAxisDescriptor + from pytato.array import (EinsumElementwiseAxis, EinsumReductionAxis, + EinsumAxisDescriptor) if not isinstance(expr, Einsum): raise TypeError(f"{expr} expected to be Einsum, got {type(expr)}.") @@ -243,7 +244,7 @@ def is_einsum_similar_to_subscript(expr: Einsum, subscripts: str) -> bool: for idim, idx in enumerate(_get_indices_from_input_subscript(out_spec, is_output=True)): - index_to_descrs[idx] = ElementwiseAxis(idim) + index_to_descrs[idx] = EinsumElementwiseAxis(idim) if len(in_spec.split(",")) != len(expr.args): return False @@ -263,7 +264,7 @@ def is_einsum_similar_to_subscript(expr: Einsum, subscripts: str) -> bool: if index_to_descrs[idx] != access_descr: return False except KeyError: - if not isinstance(access_descr, ReductionAxis): + if not isinstance(access_descr, EinsumReductionAxis): return False index_to_descrs[idx] = access_descr diff --git a/pytato/array.py b/pytato/array.py index 03f6f4d35..45cdfb7e7 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -137,8 +137,8 @@ ------------ .. autoclass:: EinsumAxisDescriptor -.. autoclass:: ElementwiseAxis -.. autoclass:: ReductionAxis +.. autoclass:: EinsumElementwiseAxis +.. autoclass:: EinsumReductionAxis .. autoclass:: NormalizedSlice Internal stuff that is only here because the documentation tool wants it @@ -858,20 +858,20 @@ class EinsumAxisDescriptor: @dataclass(eq=True, frozen=True) -class ElementwiseAxis(EinsumAxisDescriptor): +class EinsumElementwiseAxis(EinsumAxisDescriptor): """ Describes an elementwise access pattern of an array's axis. In terms of the - nomenclature used by :class:`IndexLambda`, ``ElementwiseAxis(dim=1)`` would + nomenclature used by :class:`IndexLambda`, ``EinsumElementwiseAxis(dim=1)`` would correspond to indexing the array's axis as ``_1`` in the expression. """ dim: int @dataclass(eq=True, frozen=True) -class ReductionAxis(EinsumAxisDescriptor): +class EinsumReductionAxis(EinsumAxisDescriptor): """ Describes a reduction access pattern of an array's axis. In terms of the - nomenclature used by :class:`IndexLambda`, ``ReductionAxis(dim=0)`` would + nomenclature used by :class:`IndexLambda`, ``EinsumReductionAxis(dim=0)`` would correspond to indexing the array's axis as ``_r0`` in the expression. """ dim: int @@ -945,9 +945,9 @@ def shape(self) -> ShapeType: iaxis_to_len: Dict[int, ShapeComponent] = {} for descr, axis_len in self._access_descr_to_axis_len().items(): - if isinstance(descr, ElementwiseAxis): + if isinstance(descr, EinsumElementwiseAxis): iaxis_to_len[descr.dim] = axis_len - elif isinstance(descr, ReductionAxis): + elif isinstance(descr, EinsumReductionAxis): # reduction axes do not count towards einsum's shape pass else: @@ -972,7 +972,7 @@ def _normalize_einsum_out_subscript(subscript: str) -> PMap[str, """ Normalizes the output subscript of an einsum (provided in the explicit mode). Returns a mapping from index name to an instance of - :class:`ElementwiseAxis`. + :class:`EinsumElementwiseAxis`. .. testsetup:: @@ -984,8 +984,8 @@ def _normalize_einsum_out_subscript(subscript: str) -> PMap[str, >>> sorted(result.keys()) ['i', 'j', 'k'] >>> result["i"], result["j"], result["k"] - (ElementwiseAxis(dim=1), ElementwiseAxis(dim=2), ElementwiseAxis(dim=0)) - """ + (EinsumElementwiseAxis(dim=1), EinsumElementwiseAxis(dim=2), EinsumElementwiseAxis(dim=0)) + """ # noqa: E501 normalized_indices: List[str] = [] acc = subscript.strip() @@ -1007,7 +1007,7 @@ def _normalize_einsum_out_subscript(subscript: str) -> PMap[str, raise ValueError("Used an input more than once to refer to the" f" output axis in '{subscript}") - return pmap({idx: ElementwiseAxis(i) + return pmap({idx: EinsumElementwiseAxis(i) for i, idx in enumerate(normalized_indices)}) @@ -1086,8 +1086,8 @@ def _normalize_einsum_in_subscript(subscript: str, in_axis_len) else: redn_sr_no = len([descr for descr in index_to_descr.values() - if isinstance(descr, ReductionAxis)]) - redn_axis_descr = ReductionAxis(redn_sr_no) + if isinstance(descr, EinsumReductionAxis)]) + redn_axis_descr = EinsumReductionAxis(redn_sr_no) index_to_descr = index_to_descr.set(index_char, redn_axis_descr) index_to_axis_length = index_to_axis_length.set(index_char, in_axis_len) @@ -1139,7 +1139,7 @@ def einsum(subscripts: str, *operands: Array) -> Einsum: axes=_get_default_axes(len({descr for descr in index_to_descr.values() if isinstance(descr, - ElementwiseAxis)}) + EinsumElementwiseAxis)}) )) # }}} diff --git a/pytato/codegen.py b/pytato/codegen.py index 653c2a2e1..56159af39 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -287,7 +287,7 @@ def map_einsum(self, expr: Einsum) -> Array: from pytato.scalar_expr import Reduce from pytato.utils import (dim_to_index_lambda_components, are_shape_components_equal) - from pytato.array import ElementwiseAxis, ReductionAxis + from pytato.array import EinsumElementwiseAxis, EinsumReductionAxis bindings = {f"in{k}": self.rec(arg) for k, arg in enumerate(expr.args)} redn_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] = {} @@ -308,10 +308,10 @@ def map_einsum(self, expr: Einsum) -> Array: subscript_indices.append(0) continue - if isinstance(axis, ElementwiseAxis): + if isinstance(axis, EinsumElementwiseAxis): subscript_indices.append(prim.Variable(f"_{axis.dim}")) else: - assert isinstance(axis, ReductionAxis) + assert isinstance(axis, EinsumReductionAxis) redn_idx_name = f"_r{axis.dim}" if redn_idx_name not in redn_bounds: # convert the ShapeComponent to a ScalarExpression diff --git a/pytato/target/python/numpy_like.py b/pytato/target/python/numpy_like.py index 080485c49..4451af262 100644 --- a/pytato/target/python/numpy_like.py +++ b/pytato/target/python/numpy_like.py @@ -96,7 +96,7 @@ def first_true(iterable: Iterable[T], default: T, def _get_einsum_subscripts(einsum: Einsum) -> str: - from pytato.array import (ElementwiseAxis, EinsumAxisDescriptor) + from pytato.array import EinsumElementwiseAxis, EinsumAxisDescriptor idx_stream = (chr(i) for i in range(ord("i"), ord("z"))) idx_gen: Callable[[], str] = lambda: next(idx_stream) # noqa: E731 @@ -113,7 +113,7 @@ def _get_einsum_subscripts(einsum: Einsum) -> str: input_specs.append(spec) - output_spec = "".join(axis_descr_to_idx[ElementwiseAxis(i)] + output_spec = "".join(axis_descr_to_idx[EinsumElementwiseAxis(i)] for i in range(einsum.ndim)) return f"{', '.join(input_specs)} -> {output_spec}" From 42cade58c71ed3d97a10358e57567777dfe18174 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Fri, 3 Dec 2021 18:29:34 +0530 Subject: [PATCH 02/10] adds helpers to get all reduction induction variables --- pytato/scalar_expr.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index a84ca09aa..d1767e117 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -26,7 +26,8 @@ from numbers import Number from typing import ( - Any, Union, Mapping, FrozenSet, Set, Tuple, Optional, TYPE_CHECKING) + Any, Union, Mapping, FrozenSet, Set, Tuple, Optional, TYPE_CHECKING, + Iterable) from pymbolic.mapper import (WalkMapper as WalkMapperBase, IdentityMapper as IdentityMapperBase) @@ -40,6 +41,7 @@ DistributeMapperBase) from pymbolic.mapper.stringifier import (StringifyMapper as StringifyMapperBase) +from pymbolic.mapper import CombineMapper as CombineMapperBase from pymbolic.mapper.collector import TermCollector as TermCollectorBase import pymbolic.primitives as prim import numpy as np @@ -98,6 +100,13 @@ def map_reduce(self, expr: Reduce) -> None: self.post_visit(expr) +class CombineMapper(CombineMapperBase): + def map_reduce(self, expr: Reduce, *args: Any, **kwargs: Any) -> Any: + return self.combine([*(self.rec(bnd, *args, **kwargs) + for _, bnd in sorted(expr.bounds.items())), + self.rec(expr.inner_expr, *args, **kwargs)]) + + class IdentityMapper(IdentityMapperBase): pass @@ -268,4 +277,27 @@ def __getinitargs__(self) -> Tuple[ScalarExpression, ReductionOperation, Any]: # }}} + +class InductionVariableCollector(CombineMapper): + def combine(self, values: Iterable[FrozenSet[str]]) -> FrozenSet[str]: + from functools import reduce + return reduce(frozenset.union, values, frozenset()) + + def map_reduce(self, expr: Reduce) -> FrozenSet[str]: # type: ignore[override] + return self.combine([frozenset(expr.bounds.keys()), + super().map_reduce(expr)]) + + def map_algebraic_leaf(self, expr: prim.Expression) -> FrozenSet[str]: + return frozenset() + + def map_constant(self, expr: Any) -> FrozenSet[str]: + return frozenset() + + +def get_reduction_induction_variables(expr: prim.Expression) -> FrozenSet[str]: + """ + Returns the induction variables for the reduction nodes. + """ + return InductionVariableCollector()(expr) # type: ignore[no-any-return] + # vim: foldmethod=marker From 71bb7c29ac32ad55ca44a85f6434ed776b930205 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Fri, 3 Dec 2021 18:19:28 +0530 Subject: [PATCH 03/10] implements ReductionDescriptor --- pytato/__init__.py | 4 ++-- pytato/array.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/pytato/__init__.py b/pytato/__init__.py index 7f2185b30..a955df64e 100644 --- a/pytato/__init__.py +++ b/pytato/__init__.py @@ -54,7 +54,7 @@ def set_debug_enabled(flag: bool) -> None: Einsum, Stack, Concatenate, AxisPermutation, IndexBase, Roll, IndexRemappingBase, BasicIndex, AdvancedIndexInContiguousAxes, AdvancedIndexInNoncontiguousAxes, - SizeParam, Axis, + SizeParam, Axis, ReductionDescriptor, make_dict_of_named_arrays, make_placeholder, make_size_param, make_data_wrapper, @@ -113,7 +113,7 @@ def set_debug_enabled(flag: bool) -> None: "Stack", "Concatenate", "AxisPermutation", "IndexBase", "Roll", "IndexRemappingBase", "AdvancedIndexInContiguousAxes", "AdvancedIndexInNoncontiguousAxes", - "BasicIndex", "SizeParam", "Axis", + "BasicIndex", "SizeParam", "Axis", "ReductionDescriptor", "make_dict_of_named_arrays", "make_placeholder", "make_size_param", "make_data_wrapper", "einsum", diff --git a/pytato/array.py b/pytato/array.py index 45cdfb7e7..d51fca81a 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -42,6 +42,7 @@ .. autoclass:: Array .. autoclass:: Axis +.. autoclass:: ReductionDescriptor .. autoclass:: NamedArray .. autoclass:: DictOfNamedArrays .. autoclass:: AbstractResultWithNamedArrays @@ -322,6 +323,19 @@ def _with_new_tags(self, tags: FrozenSet[Tag]) -> Taggable: return replace(self, tags=tags) +@dataclass(eq=True, frozen=True) +class ReductionDescriptor(Taggable): + """ + Records information about a reduction dimension in an + :class:`~pytato.Array`'. + """ + tags: FrozenSet[Tag] + + def _with_new_tags(self, tags: FrozenSet[Tag]) -> ReductionDescriptor: + from dataclasses import replace + return replace(self, tags=tags) + + class Array(Taggable): r""" A base class (abstract interface + supplemental functionality) for lazily From f550904836efa469d97e1276acf7119d8bf9942d Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Fri, 3 Dec 2021 18:29:00 +0530 Subject: [PATCH 04/10] defines IndexLambda.var_to_reduction_descr --- pytato/array.py | 102 +++++++++++++++++++--- pytato/cmath.py | 5 +- pytato/codegen.py | 8 ++ pytato/equality.py | 4 +- pytato/reductions.py | 162 ++++++++++++++++++++++++++++------- pytato/stringifier.py | 14 ++- pytato/transform/__init__.py | 4 + pytato/utils.py | 2 + 8 files changed, 254 insertions(+), 47 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index d51fca81a..c465fa687 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -169,9 +169,10 @@ from pytato.scalar_expr import (ScalarType, SCALAR_CLASSES, ScalarExpression, IntegralT, - INT_CLASSES) + INT_CLASSES, get_reduction_induction_variables) import re -from pyrsistent import pmap, PMap +from pyrsistent import pmap +from pyrsistent.typing import PMap # {{{ get a type variable that represents the type of '...' @@ -555,7 +556,8 @@ def _unary_op(self, op: Any) -> Array: dtype=self.dtype, bindings=bindings, tags=_get_default_tags(), - axes=_get_default_axes(self.ndim)) + axes=_get_default_axes(self.ndim), + var_to_reduction_descr=pmap()) __mul__ = partialmethod(_binary_op, operator.mul) __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) @@ -840,9 +842,16 @@ class IndexLambda(_SuppliedShapeAndDtypeMixin, Array): expressions available for use in :attr:`expr`. + .. attribute:: var_to_reduction_descr + + A mapping from reduction variables in :attr:`expr` to their + :class:`ReductionDescriptor`. + + .. automethod:: with_tagged_reduction """ - _fields = Array._fields + ("expr", "shape", "dtype", "bindings") + _fields = Array._fields + ("expr", "shape", "dtype", + "bindings", "var_to_reduction_descr") _mapper_method = "map_index_lambda" def __init__(self, @@ -851,12 +860,50 @@ def __init__(self, dtype: np.dtype[Any], bindings: Dict[str, Array], axes: AxesT, + var_to_reduction_descr: PMap[str, ReductionDescriptor], tags: FrozenSet[Tag] = frozenset()): super().__init__(shape=shape, dtype=dtype, axes=axes, tags=tags) self.expr = expr self.bindings = bindings + self.var_to_reduction_descr = var_to_reduction_descr + + def with_tagged_reduction(self, + reduction_variable: str, + tag: Tag) -> IndexLambda: + """ + Returns a copy of *self* with the :class:`ReductionDescriptor` + associated with *reduction_variable* tagged with *tag*. + + :arg reduction_variable: Name of reduction variable in *self* that + is to be tagged. + """ + from pytato.diagnostic import NotAReductionAxis + if not isinstance(reduction_variable, str): + raise TypeError("Argument 'reduction_variable' expected to be str, " + f"got {type(reduction_variable)}.") + + assert (frozenset(self.var_to_reduction_descr) + == get_reduction_induction_variables(self.expr)) + + if reduction_variable not in self.var_to_reduction_descr: + raise NotAReductionAxis( + "reduction_variable can be one of" + f" '{self.var_to_reduction_descr.keys()}'," + f" got '{reduction_variable}'.") + + new_var_to_redn_descr = self.var_to_reduction_descr.set( + reduction_variable, + self.var_to_reduction_descr[reduction_variable].tagged(tag)) + + return type(self)(expr=self.expr, + shape=self.shape, + dtype=self.dtype, + bindings=self.bindings, + axes=self.axes, + var_to_reduction_descr=new_var_to_redn_descr, + tags=self.tags) # }}} @@ -2022,7 +2069,8 @@ def full(shape: ConvertibleToShape, fill_value: ScalarType, return IndexLambda(fill_value, shape, dtype, {}, tags=_get_default_tags(), - axes=_get_default_axes(len(shape))) + axes=_get_default_axes(len(shape)), + var_to_reduction_descr=pmap()) def zeros(shape: ConvertibleToShape, dtype: Any = float, @@ -2067,7 +2115,8 @@ def eye(N: int, M: Optional[int] = None, k: int = 0, # noqa: N803 return IndexLambda(parse(f"1 if ((_1 - _0) == {k}) else 0"), shape=(N, M), dtype=dtype, bindings={}, tags=_get_default_tags(), - axes=_get_default_axes(2)) + axes=_get_default_axes(2), + var_to_reduction_descr=pmap()) # }}} @@ -2161,7 +2210,8 @@ def arange(*args: Any, **kwargs: Any) -> Array: return IndexLambda(start + Variable("_0") * step, shape=(size,), dtype=dtype, bindings={}, tags=_get_default_tags(), - axes=_get_default_axes(1)) + axes=_get_default_axes(1), + var_to_reduction_descr=pmap()) # }}} @@ -2262,7 +2312,8 @@ def logical_not(x: ArrayOrScalar) -> Union[Array, bool]: dtype=np.dtype(np.bool8), bindings={"_in0": x}, tags=_get_default_tags(), - axes=_get_default_axes(len(x.shape))) + axes=_get_default_axes(len(x.shape)), + var_to_reduction_descr=pmap()) # }}} @@ -2315,7 +2366,8 @@ def where(condition: ArrayOrScalar, dtype=dtype, bindings=bindings, tags=_get_default_tags(), - axes=_get_default_axes(len(result_shape))) + axes=_get_default_axes(len(result_shape)), + var_to_reduction_descr=pmap()) # }}} @@ -2368,10 +2420,15 @@ def make_index_lambda( expression: Union[str, ScalarExpression], bindings: Dict[str, Array], shape: ShapeType, - dtype: Any) -> IndexLambda: + dtype: Any, + var_to_reduction_descr: Optional[Mapping[str, ReductionDescriptor]] = None +) -> IndexLambda: if isinstance(expression, str): raise NotImplementedError + if var_to_reduction_descr is None: + var_to_reduction_descr = {} + # {{{ sanity checks from pytato.scalar_expr import get_dependencies @@ -2383,12 +2440,32 @@ def make_index_lambda( # }}} + # {{{ process var_to_reduction_descr + + processed_var_to_reduction_descr = {} + redn_vars = get_reduction_induction_variables(expression) + + if not (frozenset(var_to_reduction_descr) <= redn_vars): + raise ValueError(f"'{frozenset(var_to_reduction_descr) - redn_vars}': not" + " reduction induction variables.") + + for redn_var in redn_vars: + redn_descr = var_to_reduction_descr.get(redn_var, + ReductionDescriptor(frozenset())) + if not isinstance(redn_descr, ReductionDescriptor): + raise TypeError(f"reduction_dim for {redn_var} expected to be" + f" of type ReductionDescriptor, got {type(redn_descr)}.") + processed_var_to_reduction_descr[redn_var] = redn_descr + + # }}} + return IndexLambda(expr=expression, bindings=bindings, shape=shape, dtype=dtype, tags=_get_default_tags(), - axes=_get_default_axes(len(shape))) + axes=_get_default_axes(len(shape)), + var_to_reduction_descr=pmap(processed_var_to_reduction_descr)) # }}} @@ -2466,7 +2543,8 @@ def broadcast_to(array: Array, shape: ShapeType) -> Array: dtype=array.dtype, bindings={"in": array}, tags=_get_default_tags(), - axes=_get_default_axes(len(shape))) + axes=_get_default_axes(len(shape)), + var_to_reduction_descr=pmap()) def squeeze(array: Array) -> Array: diff --git a/pytato/cmath.py b/pytato/cmath.py index 5ada73e6b..06ae2c175 100644 --- a/pytato/cmath.py +++ b/pytato/cmath.py @@ -62,6 +62,7 @@ _get_default_axes, _get_default_tags) from pytato.scalar_expr import SCALAR_CLASSES from pymbolic import var +from pyrsistent import pmap def _apply_elem_wise_func(inputs: Tuple[ArrayOrScalar, ...], @@ -114,7 +115,9 @@ def _apply_elem_wise_func(inputs: Tuple[ArrayOrScalar, ...], tuple(sym_args)), shape, ret_dtype, bindings, tags=_get_default_tags(), - axes=_get_default_axes(len(shape))) + axes=_get_default_axes(len(shape)), + var_to_reduction_descr=pmap(), + ) def abs(x: Array) -> ArrayOrScalar: diff --git a/pytato/codegen.py b/pytato/codegen.py index 56159af39..f20aaddae 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -45,6 +45,7 @@ from pytato.loopy import LoopyCall from pytato.tags import AssumeNonNegative from pytools import UniqueNameGenerator +from pyrsistent import pmap import loopy as lp SymbolicIndex = Tuple[IntegralScalarExpression, ...] @@ -206,6 +207,7 @@ def get_subscript(array_index: int) -> SymbolicIndex: dtype=expr.dtype, axes=expr.axes, bindings=bindings, + var_to_reduction_descr=pmap(), tags=expr.tags) def map_concatenate(self, expr: Concatenate) -> Array: @@ -253,6 +255,7 @@ def get_subscript(array_index: int, offset: ScalarExpression) -> Subscript: dtype=expr.dtype, bindings=bindings, axes=expr.axes, + var_to_reduction_descr=pmap(), tags=expr.tags) def map_roll(self, expr: Roll) -> Array: @@ -279,6 +282,7 @@ def map_roll(self, expr: Roll) -> Array: bindings={name: self.rec(bnd) for name, bnd in bindings.items()}, axes=expr.axes, + var_to_reduction_descr=pmap(), tags=expr.tags) def map_einsum(self, expr: Einsum) -> Array: @@ -366,6 +370,7 @@ def handle_index_remapping(self, dtype=expr.dtype, bindings=dict(_in0=array), axes=expr.axes, + var_to_reduction_descr=pmap(), tags=expr.tags) def _indices_for_axis_permutation(self, expr: AxisPermutation) -> SymbolicIndex: @@ -438,6 +443,7 @@ def map_basic_index(self, expr: BasicIndex) -> IndexLambda: shape=expr.shape, dtype=expr.dtype, axes=expr.axes, + var_to_reduction_descr=pmap(), tags=expr.tags, ) @@ -503,6 +509,7 @@ def map_contiguous_advanced_index(self, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, + var_to_reduction_descr=pmap(), tags=expr.tags, ) @@ -565,6 +572,7 @@ def map_non_contiguous_advanced_index(self, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, + var_to_reduction_descr=pmap(), tags=expr.tags, ) # }}} diff --git a/pytato/equality.py b/pytato/equality.py index f864769c9..7f2cc09ab 100644 --- a/pytato/equality.py +++ b/pytato/equality.py @@ -130,7 +130,9 @@ def map_index_lambda(self, expr1: IndexLambda, expr2: Any) -> bool: else dim1 == dim2 for dim1, dim2 in zip(expr1.shape, expr2.shape)) and expr1.tags == expr2.tags - and expr1.axes == expr2.axes) + and expr1.axes == expr2.axes + and expr1.var_to_reduction_descr == expr2.var_to_reduction_descr + ) def map_stack(self, expr1: Stack, expr2: Any) -> bool: return (expr1.__class__ is expr2.__class__ diff --git a/pytato/reductions.py b/pytato/reductions.py index 04212c0a1..1213b6ec4 100644 --- a/pytato/reductions.py +++ b/pytato/reductions.py @@ -27,13 +27,14 @@ THE SOFTWARE. """ -from typing import Any, Optional, Tuple, Union, Sequence, Dict, List +from typing import Any, Optional, Tuple, Union, Sequence, Dict, List, Mapping from abc import ABC, abstractmethod import numpy as np -from pytato.array import ShapeType, Array, make_index_lambda +from pytato.array import ShapeType, Array, make_index_lambda, ReductionDescriptor from pytato.scalar_expr import ScalarExpression, Reduce, INT_CLASSES +from pyrsistent import pmap, PMap import pymbolic.primitives as prim # {{{ docs @@ -142,7 +143,7 @@ def neutral_element(self, dtype: np.dtype[Any]) -> Any: def _normalize_reduction_axes( shape: ShapeType, - reduction_axes: Optional[Union[int, Tuple[int]]] + reduction_axes: Optional[Union[int, Tuple[int, ...]]] ) -> Tuple[ShapeType, Tuple[int, ...]]: """ Returns a :class:`tuple` of ``(new_shape, normalized_redn_axes)``, where @@ -178,15 +179,18 @@ def _normalize_reduction_axes( def _get_reduction_indices_bounds(shape: ShapeType, - axes: Tuple[int, ...]) -> Tuple[ - Sequence[ScalarExpression], - Dict[str, Tuple[ScalarExpression, ScalarExpression]]]: - """Given *shape* and reduction axes *axes*, produce a list of inames + axes: Tuple[int, ...], + ) -> Tuple[Sequence[prim.Variable], + PMap[str, Tuple[ScalarExpression, + ScalarExpression]]]: + """ + Given *shape* and reduction axes *axes*, produce a list of inames ``indices`` named appropriately for reduction inames. Also fill a dictionary with bounds for reduction inames ``redn_bounds = {red_iname: (lower_bound, upper_bound)}``, where the bounds are given as a Python-style half-open interval. - :returns: ``indices, redn_bounds`` + + :returns: ``indices, redn_bounds, var_to_redn_descr`` """ indices: List[prim.Variable] = [] redn_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] = {} @@ -203,20 +207,65 @@ def _get_reduction_indices_bounds(shape: ShapeType, idx = f"_r{n_redn_dims}" indices.append(prim.Variable(idx)) redn_bounds[idx] = (0, axis_len) + n_redn_dims += 1 else: indices.append(prim.Variable(f"_{n_out_dims}")) n_out_dims += 1 - from pyrsistent import pmap + return indices, pmap(redn_bounds) + + +def _get_var_to_redn_descr(shape: ShapeType, + axes: Tuple[int, ...], + axis_to_reduction_descr: Optional[ + Mapping[int, + ReductionDescriptor]] + ) -> PMap[str, ReductionDescriptor]: + """ + :arg axis_to_reduction_descr: Mapping from a reduction axis to + its instance of :class:`~pytato.ReductionDescriptor`. This mapping + is provided by the caller of top-level functions like + :func:`pytato.sum`, :func:`pytato.prod`. + """ + var_to_redn_descr = {} + + if axis_to_reduction_descr is None: + axis_to_reduction_descr = {} + + if not (frozenset(axis_to_reduction_descr) <= frozenset(axes)): + raise ValueError("Axes " + f"'{frozenset(axis_to_reduction_descr) - frozenset(axes)}'" + " in 'axis_to_reduction_descr' not a part of axes" + " to be reduced over.") + + n_redn_dims = 0 + for idim, axis_len in enumerate(shape): + if idim in axes: + if not isinstance(axis_len, INT_CLASSES): + # TODO: add bindings for shape array expressions + raise NotImplementedError("Parametric shapes for reduction axes" + " not yet supported.") + + idx = f"_r{n_redn_dims}" + redn_descr = axis_to_reduction_descr.get( + idim, + ReductionDescriptor(frozenset())) + if not isinstance(redn_descr, ReductionDescriptor): + raise TypeError(f"'axis_to_reduction_descr[{idim}]': " + "expected an instance of ReductionDescriptor, " + f"got {type(redn_descr)}.") + var_to_redn_descr[idx] = redn_descr + n_redn_dims += 1 - # insufficient type annotation in pyrsistent - return indices, pmap(redn_bounds) # type: ignore + return pmap(var_to_redn_descr) -def _make_reduction_lambda(op: ReductionOperation, a: Array, - axis: Optional[Union[int, Tuple[int]]], - initial: Any) -> Array: +def _make_reduction_lambda( + op: ReductionOperation, a: Array, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None, + initial: Any = _NoValue) -> Array: """ Return a :class:`IndexLambda` that performs reduction over the *axis* axes of *a* with the reduction op *op*. @@ -230,7 +279,12 @@ def _make_reduction_lambda(op: ReductionOperation, a: Array, """ new_shape, reduction_axes = _normalize_reduction_axes(a.shape, axis) del axis - indices, redn_bounds = _get_reduction_indices_bounds(a.shape, reduction_axes) + indices, redn_bounds = _get_reduction_indices_bounds(a.shape, + reduction_axes) + + var_to_redn_descr = _get_var_to_redn_descr(a.shape, + reduction_axes, + axis_to_reduction_descr) if initial is _NoValue: for iax in reduction_axes: @@ -258,11 +312,15 @@ def _make_reduction_lambda(op: ReductionOperation, a: Array, redn_bounds), {"in": a}, new_shape, - a.dtype) + a.dtype, + var_to_reduction_descr=var_to_redn_descr) -def sum(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, - initial: Any = 0) -> Array: +def sum(a: Array, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + initial: Any = 0, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None + ) -> Array: """ Sums array *a*'s elements along the *axis* axes. @@ -273,12 +331,19 @@ def sum(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, :arg 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. + + :arg axis_to_reduction_descr: A mapping from axis in *axis* to the + corresponding instance of :class:`~pytato.ReductionDescriptor` that the + :class:`~pytato.array.IndexLambda` is to be instantiated with. """ - return _make_reduction_lambda(SumReductionOperation(), a, axis, initial) + return _make_reduction_lambda(SumReductionOperation(), a, axis, + axis_to_reduction_descr, initial) def amax(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, *, - initial: Any = _NoValue) -> Array: + initial: Any = _NoValue, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None + ) -> Array: """ Returns the max of array *a*'s elements along the *axis* axes. @@ -292,12 +357,20 @@ def amax(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, *, If not supplied, an :exc:`ValueError` will be raised if the reduction is empty. In that case, the reduction size must not be symbolic. + + :arg axis_to_reduction_descr: A mapping from axis in *axis* to the + corresponding instance of :class:`~pytato.ReductionDescriptor` that the + :class:`~pytato.array.IndexLambda` is to be instantiated with. """ - return _make_reduction_lambda(MaxReductionOperation(), a, axis, initial) + return _make_reduction_lambda(MaxReductionOperation(), a, axis, + axis_to_reduction_descr, initial) -def amin(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, - initial: Any = _NoValue) -> Array: +def amin(a: Array, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + initial: Any = _NoValue, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None + ) -> Array: """ Returns the min of array *a*'s elements along the *axis* axes. @@ -311,12 +384,19 @@ def amin(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, If not supplied, an :exc:`ValueError` will be raised if the reduction is empty. In that case, the reduction size must not be symbolic. + :arg axis_to_reduction_descr: A mapping from axis in *axis* to the + corresponding instance of :class:`~pytato.ReductionDescriptor` that the + :class:`~pytato.array.IndexLambda` is to be instantiated with. """ - return _make_reduction_lambda(MinReductionOperation(), a, axis, initial) + return _make_reduction_lambda(MinReductionOperation(), a, axis, + axis_to_reduction_descr, initial) -def prod(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, - initial: Any = 1) -> Array: +def prod(a: Array, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + initial: Any = 1, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None + ) -> Array: """ Returns the product of array *a*'s elements along the *axis* axes. @@ -327,11 +407,18 @@ def prod(a: Array, axis: Optional[Union[int, Tuple[int]]] = None, :arg 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. + :arg axis_to_reduction_descr: A mapping from axis in *axis* to the + corresponding instance of :class:`~pytato.ReductionDescriptor` that the + :class:`~pytato.array.IndexLambda` is to be instantiated with. """ - return _make_reduction_lambda(ProductReductionOperation(), a, axis, initial) + return _make_reduction_lambda(ProductReductionOperation(), a, axis, + axis_to_reduction_descr, initial) -def all(a: Array, axis: Optional[Union[int, Tuple[int]]] = None) -> Array: +def all(a: Array, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None + ) -> Array: """ Returns the logical-and array *a*'s elements along the *axis* axes. @@ -339,11 +426,19 @@ def all(a: Array, axis: Optional[Union[int, Tuple[int]]] = None) -> Array: :arg axis: The axes along which the elements are to be product-reduced. Defaults to all axes of the input array. + + :arg axis_to_reduction_descr: A mapping from axis in *axis* to the + corresponding instance of :class:`~pytato.ReductionDescriptor` that the + :class:`~pytato.array.IndexLambda` is to be instantiated with. """ - return _make_reduction_lambda(AllReductionOperation(), a, axis, initial=True) + return _make_reduction_lambda(AllReductionOperation(), a, axis, + axis_to_reduction_descr, initial=True) -def any(a: Array, axis: Optional[Union[int, Tuple[int]]] = None) -> Array: +def any(a: Array, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + axis_to_reduction_descr: Optional[Mapping[int, ReductionDescriptor]] = None + ) -> Array: """ Returns the logical-or of array *a*'s elements along the *axis* axes. @@ -351,8 +446,13 @@ def any(a: Array, axis: Optional[Union[int, Tuple[int]]] = None) -> Array: :arg axis: The axes along which the elements are to be product-reduced. Defaults to all axes of the input array. + + :arg axis_to_reduction_descr: A mapping from axis in *axis* to the + corresponding instance of :class:`~pytato.ReductionDescriptor` that the + :class:`~pytato.array.IndexLambda` is to be instantiated with. """ - return _make_reduction_lambda(AnyReductionOperation(), a, axis, initial=False) + return _make_reduction_lambda(AnyReductionOperation(), a, axis, + axis_to_reduction_descr, initial=False) # }}} diff --git a/pytato/stringifier.py b/pytato/stringifier.py index b2fd5bf03..452900f09 100644 --- a/pytato/stringifier.py +++ b/pytato/stringifier.py @@ -28,8 +28,10 @@ from typing import Any, Dict, Tuple from pytato.transform import Mapper -from pytato.array import Array, DataWrapper, DictOfNamedArrays, Axis +from pytato.array import (Array, DataWrapper, DictOfNamedArrays, Axis, + IndexLambda, ReductionDescriptor) from pytato.loopy import LoopyCall +from pyrsistent import PMap __doc__ = """ @@ -74,7 +76,7 @@ def __call__(self, expr: Any, depth: int = 0) -> str: # type: ignore[override] def map_foreign(self, expr: Any, depth: int) -> str: # type: ignore[override] if isinstance(expr, tuple): return "(" + ", ".join(self.rec(el, depth) for el in expr) + ")" - elif isinstance(expr, dict): + elif isinstance(expr, (dict, PMap)): return ("{" + ", ".join(f"{key!r}: {self.rec(val, depth)}" for key, val in expr.items()) @@ -106,6 +108,14 @@ def _map_generic_array(self, expr: Array, depth: int) -> str: # prettify: if trivial 'expr.axes' => don't print. fields = tuple(field for field in fields if field != "axes") + if (isinstance(expr, IndexLambda) + and all(redn_descr == ReductionDescriptor(frozenset()) + for redn_descr in expr.var_to_reduction_descr.values())): + # prettify: if trivial 'expr.var_to_reduction_descr' => don't print. + fields = tuple(field + for field in fields + if field != "var_to_reduction_descr") + return (f"{type(expr).__name__}(" + ", ".join(f"{field}=" f"{self.rec(getattr(expr, field), depth+1)}" diff --git a/pytato/transform/__init__.py b/pytato/transform/__init__.py index c10b80da1..04f23c1ed 100644 --- a/pytato/transform/__init__.py +++ b/pytato/transform/__init__.py @@ -223,6 +223,7 @@ def map_index_lambda(self, expr: IndexLambda) -> Array: dtype=expr.dtype, bindings=bindings, axes=expr.axes, + var_to_reduction_descr=expr.var_to_reduction_descr, tags=expr.tags) def map_placeholder(self, expr: Placeholder) -> Array: @@ -399,6 +400,7 @@ def map_index_lambda(self, expr: IndexLambda, dtype=expr.dtype, bindings=bindings, axes=expr.axes, + var_to_reduction_descr=expr.var_to_reduction_descr, tags=expr.tags) def map_placeholder(self, expr: Placeholder, *args: Any, **kwargs: Any) -> Array: @@ -1084,6 +1086,7 @@ def map_index_lambda(self, expr: IndexLambda) -> MPMSMaterializerAccumulator: {bnd_name: bnd.expr for bnd_name, bnd in children_rec.items()}, axes=expr.axes, + var_to_reduction_descr=expr.var_to_reduction_descr, tags=expr.tags) return _materialize_if_mpms(new_expr, self.nsuccessors[expr], children_rec.values()) @@ -1572,6 +1575,7 @@ def map_index_lambda(self, expr: IndexLambda, *args: Any) -> IndexLambda: bindings={name: self.handle_edge(expr, child) for name, child in sorted(expr.bindings.items())}, axes=expr.axes, + var_to_reduction_descr=expr.var_to_reduction_descr, tags=expr.tags) def map_einsum(self, expr: Einsum, *args: Any) -> Einsum: diff --git a/pytato/utils.py b/pytato/utils.py index 2c51515a8..7eb629c83 100644 --- a/pytato/utils.py +++ b/pytato/utils.py @@ -38,6 +38,7 @@ SCALAR_CLASSES, INT_CLASSES, BoolT) from pytools import UniqueNameGenerator from pytato.transform import Mapper +from pyrsistent import pmap __doc__ = """ @@ -205,6 +206,7 @@ def broadcast_binary_op(a1: ArrayOrScalar, a2: ArrayOrScalar, dtype=result_dtype, bindings=bindings, tags=_get_default_tags(), + var_to_reduction_descr=pmap(), axes=_get_default_axes(len(result_shape))) From 4d1855d306474c24caa9723ffdcc320f3367f5a2 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Fri, 3 Dec 2021 18:42:08 +0530 Subject: [PATCH 05/10] defines Einsum.redn_axis_to_redn_descr --- pytato/array.py | 94 ++++++++++++++++++++++++++++++++++-- pytato/codegen.py | 4 ++ pytato/equality.py | 1 + pytato/transform/__init__.py | 8 +++ 4 files changed, 104 insertions(+), 3 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index c465fa687..b6c006897 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -960,18 +960,35 @@ class Einsum(Array): A :class:`tuple` of array over which the Einstein summation is being performed. + + .. attribute:: access_descr_to_index + + Mapping from the access descriptors to the index used by the user during + the instantiation of the :class:`Einsum` node. This is a strictly + non-semantic attribute and only present to support a friendlier + :meth:`with_tagged_reduction`. + + .. automethod:: with_tagged_reduction """ - _fields = Array._fields + ("access_descriptors", "args") + _fields = Array._fields + ("access_descriptors", + "args", + "redn_axis_to_redn_descr", + "index_to_access_descr") _mapper_method = "map_einsum" def __init__(self, access_descriptors: Tuple[Tuple[EinsumAxisDescriptor, ...], ...], args: Tuple[Array, ...], axes: AxesT, + redn_axis_to_redn_descr: PMap[EinsumReductionAxis, + ReductionDescriptor], + index_to_access_descr: PMap[str, EinsumAxisDescriptor], tags: FrozenSet[Tag] = frozenset()): super().__init__(axes=axes, tags=tags) self.access_descriptors = access_descriptors self.args = args + self.redn_axis_to_redn_descr = redn_axis_to_redn_descr + self.index_to_access_descr = index_to_access_descr @memoize_method def _access_descr_to_axis_len(self @@ -1024,6 +1041,52 @@ def dtype(self) -> np.dtype[Any]: return np.find_common_type(array_types=[arg.dtype for arg in self.args], scalar_types=[]) + def with_tagged_reduction(self, + redn_axis: Union[EinsumReductionAxis, str], + tag: Tag) -> Einsum: + """ + Returns a copy of *self* with the :class:`ReductionDescriptor` + associated with *redn_axis* tagged with *tag*. + """ + from pytato.diagnostic import InvalidEinsumIndex, NotAReductionAxis + # {{{ sanity checks + + if isinstance(redn_axis, str): + try: + redn_axis_ = self.index_to_access_descr[redn_axis] + except KeyError: + raise InvalidEinsumIndex(f"'{redn_axis}': not a valid axis index.") + if isinstance(redn_axis_, EinsumReductionAxis): + redn_axis = redn_axis_ + else: + raise NotAReductionAxis(f"'{redn_axis}' is not" + " a reduction axis.") + elif isinstance(redn_axis, EinsumReductionAxis): + pass + else: + raise TypeError("Argument 'redn_axis' expected to be" + f" EinsumReductionAxis, got {type(redn_axis)}") + + if redn_axis in self.redn_axis_to_redn_descr: + assert any(redn_axis in access_descrs + for access_descrs in self.access_descriptors) + else: + raise ValueError(f"{redn_axis}: does not appear as a" + " reduction access descriptor.") + + # }}} + + new_redn_axis_to_redn_descr = self.redn_axis_to_redn_descr.set( + redn_axis, self.redn_axis_to_redn_descr[redn_axis].tagged(tag)) + + return type(self)(access_descriptors=self.access_descriptors, + args=self.args, + axes=self.axes, + redn_axis_to_redn_descr=new_redn_axis_to_redn_descr, + tags=self.tags, + index_to_access_descr=self.index_to_access_descr, + ) + EINSUM_FIRST_INDEX = re.compile(r"^\s*((?P[a-zA-Z])|(?P\.\.\.))\s*") @@ -1159,13 +1222,18 @@ def _normalize_einsum_in_subscript(subscript: str, index_to_descr, index_to_axis_length) -def einsum(subscripts: str, *operands: Array) -> Einsum: +def einsum(subscripts: str, *operands: Array, + index_to_redn_descr: Optional[Mapping[str, ReductionDescriptor]] = None + ) -> Einsum: """ Einstein summation *subscripts* on *operands*. """ if len(operands) == 0: raise ValueError("must specify at least one operand") + if index_to_redn_descr is None: + index_to_redn_descr = {} + if "->" not in subscripts: # implicit-mode: output spec matched by alphabetical ordering of # indices (ewwwww) @@ -1195,13 +1263,33 @@ def einsum(subscripts: str, *operands: Array) -> Einsum: index_to_axis_length)) access_descriptors.append(access_descriptor) + # {{{ process index_to_redn_descr + + redn_axis_to_redn_descr = {} + for idx, redn_descr in index_to_redn_descr.items(): + descr = index_to_descr[idx] + if isinstance(descr, EinsumReductionAxis): + redn_axis_to_redn_descr[descr] = redn_descr + else: + raise ValueError(f"'{idx}' is not a reduction dim.") + + for descr in index_to_descr.values(): + if isinstance(descr, EinsumReductionAxis): + if descr not in redn_axis_to_redn_descr: + redn_axis_to_redn_descr[descr] = ReductionDescriptor(frozenset()) + + # }}} + return Einsum(tuple(access_descriptors), operands, tags=_get_default_tags(), axes=_get_default_axes(len({descr for descr in index_to_descr.values() if isinstance(descr, EinsumElementwiseAxis)}) - )) + ), + redn_axis_to_redn_descr=pmap(redn_axis_to_redn_descr), + index_to_access_descr=index_to_descr, + ) # }}} diff --git a/pytato/codegen.py b/pytato/codegen.py index f20aaddae..ba20a73e2 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -297,6 +297,7 @@ def map_einsum(self, expr: Einsum) -> Array: redn_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] = {} args_as_pym_expr: List[prim.Subscript] = [] namegen = UniqueNameGenerator(set(bindings)) + var_to_redn_descr = {} # {{{ add bindings coming from the shape expressions @@ -326,6 +327,8 @@ def map_einsum(self, expr: Einsum) -> Array: bindings.update({k: self.rec(v) for k, v in redn_bound_bindings.items()}) + var_to_redn_descr[redn_idx_name] = ( + expr.redn_axis_to_redn_descr[axis]) subscript_indices.append(prim.Variable(redn_idx_name)) @@ -349,6 +352,7 @@ def map_einsum(self, expr: Einsum) -> Array: dtype=expr.dtype, bindings=bindings, axes=expr.axes, + var_to_reduction_descr=pmap(var_to_redn_descr), tags=expr.tags) # {{{ index remapping (roll, axis permutation, slice) diff --git a/pytato/equality.py b/pytato/equality.py index 7f2cc09ab..b8c652fce 100644 --- a/pytato/equality.py +++ b/pytato/equality.py @@ -214,6 +214,7 @@ def map_einsum(self, expr1: Einsum, expr2: Any) -> bool: expr2.args)) and expr1.tags == expr2.tags and expr1.axes == expr2.axes + and expr1.redn_axis_to_redn_descr == expr2.redn_axis_to_redn_descr ) def map_named_array(self, expr1: NamedArray, expr2: Any) -> bool: diff --git a/pytato/transform/__init__.py b/pytato/transform/__init__.py index 04f23c1ed..35e7bcbd7 100644 --- a/pytato/transform/__init__.py +++ b/pytato/transform/__init__.py @@ -290,6 +290,8 @@ def map_einsum(self, expr: Einsum) -> Array: return Einsum(expr.access_descriptors, tuple(self.rec(arg) for arg in expr.args), axes=expr.axes, + redn_axis_to_redn_descr=expr.redn_axis_to_redn_descr, + index_to_access_descr=expr.index_to_access_descr, tags=expr.tags) def map_named_array(self, expr: NamedArray) -> Array: @@ -474,6 +476,8 @@ def map_einsum(self, expr: Einsum, *args: Any, **kwargs: Any) -> Array: return Einsum(expr.access_descriptors, tuple(self.rec(arg, *args, **kwargs) for arg in expr.args), axes=expr.axes, + redn_axis_to_redn_descr=expr.redn_axis_to_redn_descr, + index_to_access_descr=expr.index_to_access_descr, tags=expr.tags) def map_named_array(self, expr: NamedArray, *args: Any, **kwargs: Any) -> Array: @@ -1163,6 +1167,8 @@ def map_einsum(self, expr: Einsum) -> MPMSMaterializerAccumulator: new_expr = Einsum(expr.access_descriptors, tuple(ary.expr for ary in rec_arrays), expr.axes, + expr.redn_axis_to_redn_descr, + expr.index_to_access_descr, expr.tags) return _materialize_if_mpms(new_expr, @@ -1584,6 +1590,8 @@ def map_einsum(self, expr: Einsum, *args: Any) -> Einsum: args=tuple(self.handle_edge(expr, arg, *args) for arg in expr.args), axes=expr.axes, + redn_axis_to_redn_descr=expr.redn_axis_to_redn_descr, + index_to_access_descr=expr.index_to_access_descr, tags=expr.tags) def map_stack(self, expr: Stack, *args: Any) -> Stack: From 5e53cad051727fba9b28209fe8c2ea512efb5e8c Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sun, 5 Dec 2021 17:58:56 +0530 Subject: [PATCH 06/10] test reduction_axis --- test/test_pytato.py | 38 ++++++++++++++++++++++++++++++++++++++ test/testlib.py | 6 ++++++ 2 files changed, 44 insertions(+) diff --git a/test/test_pytato.py b/test/test_pytato.py index 3a698c3a9..1545384db 100755 --- a/test/test_pytato.py +++ b/test/test_pytato.py @@ -901,6 +901,44 @@ def test_expand_dims_input_validate(): pt.expand_dims(a, -4) +def test_with_tagged_reduction(): + from testlib import FooRednTag + from pytato.raising import index_lambda_to_high_level_op + from pytato.diagnostic import InvalidEinsumIndex, NotAReductionAxis + x = pt.make_placeholder("x", shape=(10, 10), dtype=np.float64) + x_sum = pt.sum(x) + + with pytest.raises(NotAReductionAxis): + # axis='_0': not being reduced over. + x_sum = x_sum.with_tagged_reduction("_0", FooRednTag()) + + hlo = index_lambda_to_high_level_op(x_sum) + x_sum = x_sum.with_tagged_reduction(hlo.axes[1], FooRednTag()) + assert x_sum.var_to_reduction_descr[hlo.axes[1]].tags_of_type(FooRednTag) + assert not x_sum.var_to_reduction_descr[hlo.axes[0]].tags_of_type(FooRednTag) + + x_trace = pt.einsum("ii->i", x) + x_colsum = pt.einsum("ij->j", x) + + with pytest.raises(NotAReductionAxis): + # 'j': not being reduced over. + x_colsum.with_tagged_reduction("j", FooRednTag()) + + with pytest.raises(InvalidEinsumIndex): + # 'k': unknown axis + x_colsum.with_tagged_reduction("k", FooRednTag()) + + with pytest.raises(NotAReductionAxis): + # 'i': not being reduced over. + x_trace.with_tagged_reduction("i", FooRednTag()) + + x_colsum = x_colsum.with_tagged_reduction("i", FooRednTag()) + + assert (x_colsum + .redn_axis_to_redn_descr[x_colsum.index_to_access_descr["i"]] + .tags_of_type(FooRednTag)) + + if __name__ == "__main__": if len(sys.argv) > 1: exec(sys.argv[1]) diff --git a/test/testlib.py b/test/testlib.py index 4ff50920e..f42e1820b 100644 --- a/test/testlib.py +++ b/test/testlib.py @@ -316,6 +316,12 @@ def gen_comm(rdagc: RandomDAGContext) -> pt.Array: # {{{ tags used only by the regression tests +class FooRednTag(Tag): + """ + foo + """ + + class FooInameTag(Tag): """ foo From c4e0ce47549a017d6070f3d199905fa7e5834d8a Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Mon, 6 Dec 2021 10:38:40 +0530 Subject: [PATCH 07/10] Propagates loopy iname tags for ReductionAxis --- pytato/target/loopy/codegen.py | 35 +++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/pytato/target/loopy/codegen.py b/pytato/target/loopy/codegen.py index 43a6d48ab..81feca679 100644 --- a/pytato/target/loopy/codegen.py +++ b/pytato/target/loopy/codegen.py @@ -38,7 +38,8 @@ from pytato.array import (Array, DictOfNamedArrays, ShapeType, IndexLambda, - SizeParam, Placeholder, NamedArray, DataWrapper) + SizeParam, Placeholder, NamedArray, DataWrapper, + ReductionDescriptor) from pytato.target import BoundProgram from pytato.target.loopy import LoopyPyOpenCLTarget, LoopyTarget @@ -155,6 +156,7 @@ class LocalExpressionContext: num_indices: int local_namespace: Mapping[str, Array] reduction_bounds: ReductionBounds + var_to_reduction_descr: Mapping[str, ReductionDescriptor] def lookup(self, name: str) -> Array: return self.local_namespace[name] @@ -163,6 +165,8 @@ def copy(self, *, reduction_bounds: Optional[ReductionBounds] = None, num_indices: Optional[int] = None, local_namespace: Optional[Mapping[str, Array]] = None, + var_to_reduction_descr: Optional[ + Mapping[str, ReductionDescriptor]] = None, ) -> LocalExpressionContext: if reduction_bounds is None: reduction_bounds = self.reduction_bounds @@ -170,9 +174,12 @@ def copy(self, *, num_indices = self.num_indices if local_namespace is None: local_namespace = self.local_namespace + if var_to_reduction_descr is None: + var_to_reduction_descr = self.var_to_reduction_descr return LocalExpressionContext(reduction_bounds=reduction_bounds, num_indices=num_indices, - local_namespace=local_namespace) + local_namespace=local_namespace, + var_to_reduction_descr=var_to_reduction_descr) # }}} @@ -374,9 +381,11 @@ def map_index_lambda(self, expr: IndexLambda, return state.results[expr] prstnt_ctx = PersistentExpressionContext(state) - local_ctx = LocalExpressionContext(local_namespace=expr.bindings, - num_indices=expr.ndim, - reduction_bounds={}) + local_ctx = LocalExpressionContext( + local_namespace=expr.bindings, + num_indices=expr.ndim, + reduction_bounds={}, + var_to_reduction_descr=expr.var_to_reduction_descr) loopy_expr = self.exprgen_mapper(expr.expr, prstnt_ctx, local_ctx) result: ImplementedResult = InlinedResult(loopy_expr, @@ -507,8 +516,9 @@ def _get_sub_array_ref(array: Array, name: str) -> "lp.symbolic.SubArrayRef": depends_on.update(pt_arg_rec.depends_on) else: local_ctx = LocalExpressionContext(reduction_bounds={}, - num_indices=0, - local_namespace={}) + num_indices=0, + local_namespace={}, + var_to_reduction_descr={}) params.append(self.exprgen_mapper(pt_arg, prstnt_ctx, local_ctx)) @@ -650,6 +660,17 @@ def map_reduce(self, expr: scalar_expr.Reduce, kernel = state.kernel state.update_kernel(kernel.copy(domains=kernel.domains+[domain])) + # {{{ pytato tags -> loopy tags + + for name_in_expr, name_in_kernel in sorted(unique_names_mapping.items()): + for tag in local_ctx.var_to_reduction_descr[name_in_expr].tags: + if all(not isinstance(tag, tag_t) + for tag_t in self.codegen_mapper.axis_tag_t_to_not_propagate): + state.update_kernel(lp.tag_inames(state.kernel, + {name_in_kernel: tag})) + + # }}} + return inner_expr # }}} From 44a22c2467c029f65bedaebff8935c29f9d28789 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Mon, 6 Dec 2021 10:39:22 +0530 Subject: [PATCH 08/10] tests loopy tag propagation for ReductionAxis --- test/test_codegen.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/test/test_codegen.py b/test/test_codegen.py index a41f4fc8c..33eb4e644 100755 --- a/test/test_codegen.py +++ b/test/test_codegen.py @@ -1492,9 +1492,21 @@ def test_axis_tag_to_loopy_iname_tag_propagate(): x = pt.make_placeholder("x", (10, 4), np.float32) y = 2 * x y = (y - .with_tagged_axis(0, (FooInameTag(), BazInameTag())) - .with_tagged_axis(1, (BarInameTag(), BazInameTag()))) - t_unit = pt.generate_loopy({"y": y}, + .with_tagged_axis(0, FooInameTag()) + .with_tagged_axis(1, BarInameTag())) + x_sum = pt.sum( + x, + axis_to_reduction_descr={ + 1: pt.ReductionDescriptor(frozenset([FooInameTag()]))}) + x_einsum = pt.einsum( + "ij->", + x, + index_to_redn_descr={"i": pt.ReductionDescriptor(frozenset([BarInameTag()]))} + ) + + t_unit = pt.generate_loopy({"y": y, + "x_sum": x_sum, + "x_einsum": x_einsum}, axis_tag_t_to_not_propagate=frozenset([BazInameTag]) ).program @@ -1515,6 +1527,14 @@ def test_axis_tag_to_loopy_iname_tag_propagate(): .default_entrypoint .inames["y_dim1"] .tags_of_type(BarInameTag)) == 1 + assert len(t_unit + .default_entrypoint + .inames["_pt_sum_r1"] + .tags_of_type(FooInameTag)) == 1 + assert len(t_unit + .default_entrypoint + .inames["_pt_sum_r0_0"] + .tags_of_type(BarInameTag)) == 1 # there shouldn't be any inames tagged with BazInameTag assert len([iname From 264930ba61b3aec79baa794ae63b945e7934da40 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Wed, 9 Mar 2022 23:21:39 -0600 Subject: [PATCH 09/10] introduces ReductionAxis/EinsumIndex specific exception types --- pytato/diagnostic.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/pytato/diagnostic.py b/pytato/diagnostic.py index cfb3bccd9..6a77b32b3 100644 --- a/pytato/diagnostic.py +++ b/pytato/diagnostic.py @@ -52,3 +52,16 @@ class UnknownIndexLambdaExpr(ValueError): inferred. """ pass + + +class InvalidEinsumIndex(ValueError): + """ + Raised when an einsum index was referred by an invalid value. + """ + + +class NotAReductionAxis(ValueError): + """ + Raised when a :class:`pytato.ReductionDescriptor` was referred by an invalid + value. + """ From b7abe760b3d295c1944f7d6b6685317590c15727 Mon Sep 17 00:00:00 2001 From: Kaushik Kulkarni Date: Sun, 12 Jun 2022 12:09:35 -0500 Subject: [PATCH 10/10] prefer immutables.Map over pyrsistent.PMap --- pytato/array.py | 52 ++++++++++++++++++++++---------------------- pytato/cmath.py | 4 ++-- pytato/codegen.py | 18 +++++++-------- pytato/reductions.py | 10 ++++----- setup.py | 3 ++- 5 files changed, 44 insertions(+), 43 deletions(-) diff --git a/pytato/array.py b/pytato/array.py index b6c006897..0819dd8e2 100644 --- a/pytato/array.py +++ b/pytato/array.py @@ -171,8 +171,7 @@ ScalarExpression, IntegralT, INT_CLASSES, get_reduction_induction_variables) import re -from pyrsistent import pmap -from pyrsistent.typing import PMap +from immutables import Map # {{{ get a type variable that represents the type of '...' @@ -557,7 +556,7 @@ def _unary_op(self, op: Any) -> Array: bindings=bindings, tags=_get_default_tags(), axes=_get_default_axes(self.ndim), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) __mul__ = partialmethod(_binary_op, operator.mul) __rmul__ = partialmethod(_binary_op, operator.mul, reverse=True) @@ -860,7 +859,7 @@ def __init__(self, dtype: np.dtype[Any], bindings: Dict[str, Array], axes: AxesT, - var_to_reduction_descr: PMap[str, ReductionDescriptor], + var_to_reduction_descr: Mapping[str, ReductionDescriptor], tags: FrozenSet[Tag] = frozenset()): super().__init__(shape=shape, dtype=dtype, axes=axes, tags=tags) @@ -893,6 +892,7 @@ def with_tagged_reduction(self, f" '{self.var_to_reduction_descr.keys()}'," f" got '{reduction_variable}'.") + assert isinstance(self.var_to_reduction_descr, Map) new_var_to_redn_descr = self.var_to_reduction_descr.set( reduction_variable, self.var_to_reduction_descr[reduction_variable].tagged(tag)) @@ -980,9 +980,9 @@ def __init__(self, access_descriptors: Tuple[Tuple[EinsumAxisDescriptor, ...], ...], args: Tuple[Array, ...], axes: AxesT, - redn_axis_to_redn_descr: PMap[EinsumReductionAxis, - ReductionDescriptor], - index_to_access_descr: PMap[str, EinsumAxisDescriptor], + redn_axis_to_redn_descr: Mapping[EinsumReductionAxis, + ReductionDescriptor], + index_to_access_descr: Mapping[str, EinsumAxisDescriptor], tags: FrozenSet[Tag] = frozenset()): super().__init__(axes=axes, tags=tags) self.access_descriptors = access_descriptors @@ -992,7 +992,7 @@ def __init__(self, @memoize_method def _access_descr_to_axis_len(self - ) -> PMap[EinsumAxisDescriptor, ShapeComponent]: + ) -> Mapping[EinsumAxisDescriptor, ShapeComponent]: from pytato.utils import are_shape_components_equal descr_to_axis_len: Dict[EinsumAxisDescriptor, ShapeComponent] = {} @@ -1014,7 +1014,7 @@ def _access_descr_to_axis_len(self else: descr_to_axis_len[descr] = arg_axis_len - return pmap(descr_to_axis_len) + return Map(descr_to_axis_len) # type-ignore reason: github.com/python/mypy/issues/1362 @property # type: ignore @@ -1076,6 +1076,7 @@ def with_tagged_reduction(self, # }}} + assert isinstance(self.redn_axis_to_redn_descr, Map) new_redn_axis_to_redn_descr = self.redn_axis_to_redn_descr.set( redn_axis, self.redn_axis_to_redn_descr[redn_axis].tagged(tag)) @@ -1091,7 +1092,7 @@ def with_tagged_reduction(self, EINSUM_FIRST_INDEX = re.compile(r"^\s*((?P[a-zA-Z])|(?P\.\.\.))\s*") -def _normalize_einsum_out_subscript(subscript: str) -> PMap[str, +def _normalize_einsum_out_subscript(subscript: str) -> Map[str, EinsumAxisDescriptor]: """ Normalizes the output subscript of an einsum (provided in the explicit @@ -1131,19 +1132,19 @@ def _normalize_einsum_out_subscript(subscript: str) -> PMap[str, raise ValueError("Used an input more than once to refer to the" f" output axis in '{subscript}") - return pmap({idx: EinsumElementwiseAxis(i) + return Map({idx: EinsumElementwiseAxis(i) for i, idx in enumerate(normalized_indices)}) def _normalize_einsum_in_subscript(subscript: str, in_operand: Array, - index_to_descr: PMap[str, + index_to_descr: Map[str, EinsumAxisDescriptor], - index_to_axis_length: PMap[str, + index_to_axis_length: Map[str, ShapeComponent], ) -> Tuple[Tuple[EinsumAxisDescriptor, ...], - PMap[str, EinsumAxisDescriptor], - PMap[str, ShapeComponent]]: + Map[str, EinsumAxisDescriptor], + Map[str, ShapeComponent]]: """ Normalizes the subscript for an input operand in an einsum. Returns ``(access_descrs, updated_index_to_descr, updated_to_index_to_axis_length)``, @@ -1218,8 +1219,7 @@ def _normalize_einsum_in_subscript(subscript: str, in_operand_axis_descrs.append(index_to_descr[index_char]) - return (tuple(in_operand_axis_descrs), - index_to_descr, index_to_axis_length) + return (tuple(in_operand_axis_descrs), index_to_descr, index_to_axis_length) def einsum(subscripts: str, *operands: Array, @@ -1252,7 +1252,7 @@ def einsum(subscripts: str, *operands: Array, ) index_to_descr = _normalize_einsum_out_subscript(out_spec) - index_to_axis_length: PMap[str, ShapeComponent] = pmap() + index_to_axis_length: Map[str, ShapeComponent] = Map() access_descriptors = [] for in_spec, in_operand in zip(in_specs, operands): @@ -1287,7 +1287,7 @@ def einsum(subscripts: str, *operands: Array, if isinstance(descr, EinsumElementwiseAxis)}) ), - redn_axis_to_redn_descr=pmap(redn_axis_to_redn_descr), + redn_axis_to_redn_descr=Map(redn_axis_to_redn_descr), index_to_access_descr=index_to_descr, ) @@ -2158,7 +2158,7 @@ def full(shape: ConvertibleToShape, fill_value: ScalarType, return IndexLambda(fill_value, shape, dtype, {}, tags=_get_default_tags(), axes=_get_default_axes(len(shape)), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) def zeros(shape: ConvertibleToShape, dtype: Any = float, @@ -2204,7 +2204,7 @@ def eye(N: int, M: Optional[int] = None, k: int = 0, # noqa: N803 shape=(N, M), dtype=dtype, bindings={}, tags=_get_default_tags(), axes=_get_default_axes(2), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) # }}} @@ -2299,7 +2299,7 @@ def arange(*args: Any, **kwargs: Any) -> Array: shape=(size,), dtype=dtype, bindings={}, tags=_get_default_tags(), axes=_get_default_axes(1), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) # }}} @@ -2401,7 +2401,7 @@ def logical_not(x: ArrayOrScalar) -> Union[Array, bool]: bindings={"_in0": x}, tags=_get_default_tags(), axes=_get_default_axes(len(x.shape)), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) # }}} @@ -2455,7 +2455,7 @@ def where(condition: ArrayOrScalar, bindings=bindings, tags=_get_default_tags(), axes=_get_default_axes(len(result_shape)), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) # }}} @@ -2553,7 +2553,7 @@ def make_index_lambda( dtype=dtype, tags=_get_default_tags(), axes=_get_default_axes(len(shape)), - var_to_reduction_descr=pmap(processed_var_to_reduction_descr)) + var_to_reduction_descr=Map(processed_var_to_reduction_descr)) # }}} @@ -2632,7 +2632,7 @@ def broadcast_to(array: Array, shape: ShapeType) -> Array: bindings={"in": array}, tags=_get_default_tags(), axes=_get_default_axes(len(shape)), - var_to_reduction_descr=pmap()) + var_to_reduction_descr=Map()) def squeeze(array: Array) -> Array: diff --git a/pytato/cmath.py b/pytato/cmath.py index 06ae2c175..fdfb1654d 100644 --- a/pytato/cmath.py +++ b/pytato/cmath.py @@ -62,7 +62,7 @@ _get_default_axes, _get_default_tags) from pytato.scalar_expr import SCALAR_CLASSES from pymbolic import var -from pyrsistent import pmap +from immutables import Map def _apply_elem_wise_func(inputs: Tuple[ArrayOrScalar, ...], @@ -116,7 +116,7 @@ def _apply_elem_wise_func(inputs: Tuple[ArrayOrScalar, ...], shape, ret_dtype, bindings, tags=_get_default_tags(), axes=_get_default_axes(len(shape)), - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), ) diff --git a/pytato/codegen.py b/pytato/codegen.py index ba20a73e2..740f01a1e 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -45,7 +45,7 @@ from pytato.loopy import LoopyCall from pytato.tags import AssumeNonNegative from pytools import UniqueNameGenerator -from pyrsistent import pmap +from immutables import Map import loopy as lp SymbolicIndex = Tuple[IntegralScalarExpression, ...] @@ -207,7 +207,7 @@ def get_subscript(array_index: int) -> SymbolicIndex: dtype=expr.dtype, axes=expr.axes, bindings=bindings, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags) def map_concatenate(self, expr: Concatenate) -> Array: @@ -255,7 +255,7 @@ def get_subscript(array_index: int, offset: ScalarExpression) -> Subscript: dtype=expr.dtype, bindings=bindings, axes=expr.axes, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags) def map_roll(self, expr: Roll) -> Array: @@ -282,7 +282,7 @@ def map_roll(self, expr: Roll) -> Array: bindings={name: self.rec(bnd) for name, bnd in bindings.items()}, axes=expr.axes, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags) def map_einsum(self, expr: Einsum) -> Array: @@ -352,7 +352,7 @@ def map_einsum(self, expr: Einsum) -> Array: dtype=expr.dtype, bindings=bindings, axes=expr.axes, - var_to_reduction_descr=pmap(var_to_redn_descr), + var_to_reduction_descr=Map(var_to_redn_descr), tags=expr.tags) # {{{ index remapping (roll, axis permutation, slice) @@ -374,7 +374,7 @@ def handle_index_remapping(self, dtype=expr.dtype, bindings=dict(_in0=array), axes=expr.axes, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags) def _indices_for_axis_permutation(self, expr: AxisPermutation) -> SymbolicIndex: @@ -447,7 +447,7 @@ def map_basic_index(self, expr: BasicIndex) -> IndexLambda: shape=expr.shape, dtype=expr.dtype, axes=expr.axes, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags, ) @@ -513,7 +513,7 @@ def map_contiguous_advanced_index(self, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags, ) @@ -576,7 +576,7 @@ def map_non_contiguous_advanced_index(self, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, - var_to_reduction_descr=pmap(), + var_to_reduction_descr=Map(), tags=expr.tags, ) # }}} diff --git a/pytato/reductions.py b/pytato/reductions.py index 1213b6ec4..cd4f509f6 100644 --- a/pytato/reductions.py +++ b/pytato/reductions.py @@ -34,7 +34,7 @@ from pytato.array import ShapeType, Array, make_index_lambda, ReductionDescriptor from pytato.scalar_expr import ScalarExpression, Reduce, INT_CLASSES -from pyrsistent import pmap, PMap +from immutables import Map import pymbolic.primitives as prim # {{{ docs @@ -181,7 +181,7 @@ def _normalize_reduction_axes( def _get_reduction_indices_bounds(shape: ShapeType, axes: Tuple[int, ...], ) -> Tuple[Sequence[prim.Variable], - PMap[str, Tuple[ScalarExpression, + Mapping[str, Tuple[ScalarExpression, ScalarExpression]]]: """ Given *shape* and reduction axes *axes*, produce a list of inames @@ -213,7 +213,7 @@ def _get_reduction_indices_bounds(shape: ShapeType, indices.append(prim.Variable(f"_{n_out_dims}")) n_out_dims += 1 - return indices, pmap(redn_bounds) + return indices, Map(redn_bounds) def _get_var_to_redn_descr(shape: ShapeType, @@ -221,7 +221,7 @@ def _get_var_to_redn_descr(shape: ShapeType, axis_to_reduction_descr: Optional[ Mapping[int, ReductionDescriptor]] - ) -> PMap[str, ReductionDescriptor]: + ) -> Mapping[str, ReductionDescriptor]: """ :arg axis_to_reduction_descr: Mapping from a reduction axis to its instance of :class:`~pytato.ReductionDescriptor`. This mapping @@ -258,7 +258,7 @@ def _get_var_to_redn_descr(shape: ShapeType, var_to_redn_descr[idx] = redn_descr n_redn_dims += 1 - return pmap(var_to_redn_descr) + return Map(var_to_redn_descr) def _make_reduction_lambda( diff --git a/setup.py b/setup.py index 336157d5d..4363f3044 100644 --- a/setup.py +++ b/setup.py @@ -36,7 +36,8 @@ install_requires=[ "loopy>=2020.2", "pytools>=2021.1", - "pyrsistent" + "pyrsistent", + "immutables", ], package_data={"pytato": ["py.typed"]}, author="Andreas Kloeckner, Matt Wala, Xiaoyu Wei",