Skip to content
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

Properly handle ignored axes during tag propagation #569

Merged
merged 6 commits into from
Dec 13, 2024
Merged

Properly handle ignored axes during tag propagation #569

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0ac9864
change propagation to properly ignore ignored axes
a-alveyblanc Dec 12, 2024
2b42195
bump pytools version requirement
a-alveyblanc Dec 13, 2024
db0eb20
improve clarity of unify_axes_tags
a-alveyblanc Dec 13, 2024
52f1b3c
add test
a-alveyblanc Dec 13, 2024
ddad374
fix ruff complaints
a-alveyblanc Dec 13, 2024
871a835
actually fix ruff complaints
a-alveyblanc Dec 13, 2024
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
14 changes: 10 additions & 4 deletions pytato/transform/metadata.py
View file
Open in desktop
a-alveyblanc marked this conversation as resolved.
Show resolved Hide resolved
Original file line number Diff line number Diff line change
Expand Up @@ -716,11 +716,17 @@ def unify_axes_tags(
equations_collector.equations
)

for tag, var in equations_collector.known_tag_to_var.items():
if isinstance(tag, AxisIgnoredForPropagationTag):
continue
ignored_vars = set()
a-alveyblanc marked this conversation as resolved.
Show resolved Hide resolved
for (ary, ax), ax_var in equations_collector.axis_to_var.items():
tags = ary.axes[ax].tags_of_type(AxisIgnoredForPropagationTag)
if tags:
ignored_vars.add(ax_var)
for tag in tags:
ignored_vars.add(equations_collector.known_tag_to_var[tag])

reachable_nodes = get_reachable_nodes(propagation_graph, var)
for tag, var in equations_collector.known_tag_to_var.items():
reachable_nodes = get_reachable_nodes(propagation_graph, var,
ignored_vars)
a-alveyblanc marked this conversation as resolved.
Show resolved Hide resolved
for reachable_var in (reachable_nodes - known_tag_vars):
axis_to_solved_tags.setdefault(
equations_collector.axis_to_var.inverse[reachable_var],
Expand Down
Loading