Skip to content

Add MapAndReduceMapper and refactor/extend analysis utils #639

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

Draft
majosm wants to merge 11 commits into
base: main
Choose a base branch
from
+1,312 −410
Draft

Add MapAndReduceMapper and refactor/extend analysis utils #639

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
062154c
reference nodes from pt instead of pt.array in tests
majosm Feb 25, 2026
df55c46
add function nodes to pytato/__init__.py
majosm Feb 25, 2026
78a5037
add MapOnceMapper and MapAndReduceMapper/map_and_reduce
majosm Feb 18, 2026
dba51cc
reimplement analysis functions using MapAndReduceMapper
majosm Feb 25, 2026
a02d6d2
test counting utils for DAGs with functions
majosm Feb 25, 2026
0df4179
modify node counting tests to also count DictOfNamedArrays
majosm Feb 25, 2026
445c4a9
add get_num_node_instances_of
majosm Feb 19, 2026
6d9f119
deprecate get_num_call_sites in favor of get_num_node_instances_of
majosm Feb 19, 2026
d382451
add NodeCollector/collect_nodes
majosm Feb 19, 2026
8ba052c
add collect_node_instances_of
majosm Feb 19, 2026
0c70b89
add collect_materialized_nodes
majosm Feb 19, 2026
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
224 changes: 48 additions & 176 deletions .basedpyright/baseline.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1793,94 +1793,6 @@
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 4,
"endColumn": 7,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 11,
"endColumn": 31,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 11,
"endColumn": 31,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 4,
"endColumn": 7,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 15,
"endColumn": 35,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 15,
"endColumn": 42,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 15,
"endColumn": 44,
"lineCount": 1
}
},
{
"code": "reportUnannotatedClassAttribute",
"range": {
"startColumn": 13,
"endColumn": 18,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 4,
"endColumn": 8,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 11,
"endColumn": 21,
"lineCount": 1
}
},
{
"code": "reportAny",
"range": {
"startColumn": 11,
"endColumn": 21,
"lineCount": 1
}
},
{
"code": "reportUnannotatedClassAttribute",
"range": {
Expand Down Expand Up @@ -7513,94 +7425,6 @@
"lineCount": 1
}
},
{
"code": "reportImplicitOverride",
"range": {
"startColumn": 8,
"endColumn": 23,
"lineCount": 1
}
},
{
"code": "reportImplicitOverride",
"range": {
"startColumn": 8,
"endColumn": 24,
"lineCount": 1
}
},
{
"code": "reportImplicitOverride",
"range": {
"startColumn": 8,
"endColumn": 31,
"lineCount": 1
}
},
{
"code": "reportImplicitOverride",
"range": {
"startColumn": 8,
"endColumn": 16,
"lineCount": 1
}
},
{
"code": "reportUnusedVariable",
"range": {
"startColumn": 20,
"endColumn": 24,
"lineCount": 1
}
},
{
"code": "reportImplicitOverride",
"range": {
"startColumn": 8,
"endColumn": 15,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 15,
"endColumn": 60,
"lineCount": 1
}
},
{
"code": "reportUnknownArgumentType",
"range": {
"startColumn": 22,
"endColumn": 40,
"lineCount": 1
}
},
{
"code": "reportUnknownLambdaType",
"range": {
"startColumn": 29,
"endColumn": 30,
"lineCount": 1
}
},
{
"code": "reportUnknownLambdaType",
"range": {
"startColumn": 35,
"endColumn": 40,
"lineCount": 1
}
},
{
"code": "reportUnknownArgumentType",
"range": {
"startColumn": 48,
"endColumn": 59,
"lineCount": 1
}
},
{
"code": "reportImplicitOverride",
"range": {
Expand Down Expand Up @@ -7825,6 +7649,54 @@
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 8,
"endColumn": 14,
"lineCount": 1
}
},
{
"code": "reportUnknownMemberType",
"range": {
"startColumn": 17,
"endColumn": 27,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 15,
"endColumn": 21,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 8,
"endColumn": 14,
"lineCount": 1
}
},
{
"code": "reportUnknownMemberType",
"range": {
"startColumn": 17,
"endColumn": 47,
"lineCount": 1
}
},
{
"code": "reportUnknownVariableType",
"range": {
"startColumn": 15,
"endColumn": 21,
"lineCount": 1
}
},
{
"code": "reportIncompatibleMethodOverride",
"range": {
Expand Down
8 changes: 7 additions & 1 deletion pytato/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -149,14 +149,17 @@ def set_debug_enabled(flag: bool) -> None:
)
from pytato.distributed.tags import number_distributed_tags
from pytato.distributed.verify import verify_distributed_partition
from pytato.function import trace_call
from pytato.function import Call, FunctionDefinition, NamedCallResult, trace_call
from pytato.loopy import LoopyCall
from pytato.pad import pad
from pytato.reductions import all, amax, amin, any, prod, sum
from pytato.target import Target
from pytato.target.loopy import LoopyPyOpenCLTarget
from pytato.target.loopy.codegen import generate_loopy
from pytato.target.python.jax import generate_jax

# FIXME: Should some of the functions from pytato.transform be imported here?
# (deduplicate, map_and_copy, etc.)
from pytato.transform.calls import inline_calls, tag_all_calls_to_be_inlined
from pytato.transform.dead_code_elimination import eliminate_dead_code
from pytato.transform.lower_to_index_lambda import to_index_lambda
Expand All @@ -179,6 +182,7 @@ def set_debug_enabled(flag: bool) -> None:
"Axis",
"AxisPermutation",
"BasicIndex",
"Call",
"Concatenate",
"DataWrapper",
"DictOfNamedArrays",
Expand All @@ -188,13 +192,15 @@ def set_debug_enabled(flag: bool) -> None:
"DistributedSend",
"DistributedSendRefHolder",
"Einsum",
"FunctionDefinition",
"IndexBase",
"IndexLambda",
"IndexRemappingBase",
"InputArgumentBase",
"LoopyCall",
"LoopyPyOpenCLTarget",
"NamedArray",
"NamedCallResult",
"Placeholder",
"ReductionDescriptor",
"Reshape",
Expand Down
Loading
Loading