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.

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

Add node type counter #489

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

Conversation

MTCam
Copy link

@MTCam MTCam commented Mar 15, 2024

Adds an analysis util to count nodes of each type.

cc: @matthiasdiener

Copy link
Collaborator

@majosm majosm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks familiar. 🙂 I added something almost exactly the same in my concatenation changes.

Added a couple of suggestions here.

# {{{ NodeTypeCountMapper

@optimize_mapper(drop_args=True, drop_kwargs=True, inline_get_cache_key=True)
class NodeTypeCountMapper(CachedWalkMapper):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could maybe just modify NodeCountMapper to do this and then have get_num_nodes return sum(ncm.counts.values())? Saves some duplication.

pytato/analysis/__init__.py Outdated Show resolved Hide resolved
self.counts[type(expr)] += 1


def get_num_node_types(outputs: Union[Array, DictOfNamedArrays]) -> Dict[Type, int]:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

get_num_node_types is a little confusing (sounds like it would return the number of distinct types in the DAG). I used get_node_counts in mine, which isn't much better... get_node_type_counts? 🤷‍♂️

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed. get_node_type_counts seems like the clear choice.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants