Skip to content

Introduce attributes to store information about reduction dimensions #224

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
inducer merged 10 commits into from
Jun 13, 2022
Merged

Introduce attributes to store information about reduction dimensions #224

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
b89cf2d
Rename EinsumAxisDescriptor subclasses.
kaushikcfd Nov 30, 2021
42cade5
adds helpers to get all reduction induction variables
kaushikcfd Dec 3, 2021
71bb7c2
implements ReductionDescriptor
kaushikcfd Dec 3, 2021
f550904
defines IndexLambda.var_to_reduction_descr
kaushikcfd Dec 3, 2021
4d1855d
defines Einsum.redn_axis_to_redn_descr
kaushikcfd Dec 3, 2021
5e53cad
test reduction_axis
kaushikcfd Dec 5, 2021
c4e0ce4
Propagates loopy iname tags for ReductionAxis
kaushikcfd Dec 6, 2021
44a22c2
tests loopy tag propagation for ReductionAxis
kaushikcfd Dec 6, 2021
264930b
introduces ReductionAxis/EinsumIndex specific exception types
kaushikcfd Mar 10, 2022
b7abe76
prefer immutables.Map over pyrsistent.PMap
kaushikcfd Jun 12, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions pytato/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down
7 changes: 4 additions & 3 deletions pytato/analysis/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)}.")
Expand All @@ -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
Expand All @@ -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

Expand Down
Loading