-
Notifications
You must be signed in to change notification settings - Fork 16
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
Deterministic find_distributed_partition (non-set) #529
base: main
Are you sure you want to change the base?
Conversation
f08a309
to
3364a4f
Compare
c134b87
to
817b255
Compare
26492ce
to
f3f3c7d
Compare
This is ready for a first look @inducer. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Took a look, just a few minor nits. Looks good to go from my perspective.
pytato/distributed/partition.py
Outdated
@@ -826,7 +772,7 @@ def find_distributed_partition( | |||
raise comm_batches_or_exc | |||
|
|||
comm_batches = cast( | |||
Sequence[AbstractSet[CommunicationOpIdentifier]], | |||
Sequence[list[CommunicationOpIdentifier]], |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sequence[list[CommunicationOpIdentifier]], | |
Sequence[Collection[CommunicationOpIdentifier]], |
maybe?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I removed the cast completely in 168ef53
pytato/analysis/__init__.py
Outdated
We only consider the predecessors of a nodes in a data-flow sense. | ||
""" | ||
def _get_preds_from_shape(self, shape: ShapeType) -> frozenset[Array]: | ||
return frozenset({dim for dim in shape if isinstance(dim, Array)}) | ||
def _get_preds_from_shape(self, shape: ShapeType) -> abc_Set[Array]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
def _get_preds_from_shape(self, shape: ShapeType) -> abc_Set[Array]: | |
def _get_preds_from_shape(self, shape: ShapeType) -> AbstractSet[Array]: |
(from typing
)? I'm not sure mypy understands the collections.abc
types well.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think this may have been fixed with 076a76e
pytato/distributed/partition.py
Outdated
output_arrays = tuple(unique(outputs._data.values())) | ||
mso_arrays = tuple(unique(materialized_arrays + sent_arrays + output_arrays)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think these might be fine left as lists (as returned from) unique
, typed as Sequence
to have mypy flag attempted mutation.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not sure I understand - could you please clarify? (unique
doesn't return a list
)
3c87726
to
076a76e
Compare
191422d
to
168ef53
Compare
This is ready for another look @inducer. It seems to work fine with the |
Nevermind, this seems to have been a merge error. |
This is ready for review @inducer. The current version shows the same performance as the baseline, and I have not seen any determinism-related issues when using this PR. |
cc1045a
to
e8b5806
Compare
As far as I can see, the mypy errors are unrelated to this PR. |
@@ -327,37 +327,37 @@ class DirectPredecessorsGetter(Mapper[frozenset[ArrayOrNames], []]): | |||
|
|||
We only consider the predecessors of a nodes in a data-flow sense. | |||
""" | |||
def _get_preds_from_shape(self, shape: ShapeType) -> frozenset[ArrayOrNames]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
T = TypeVar("T")
class FakeOrderedFrozenSet(immutabledict[T, None]):
pass
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I wasn't able to get this to work with mypy, but what do you think of fbcbcef?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That works for me. I was under the impression that type aliases could not be generic, but apparently I'm wrong? I tried for a bit to back up my assumption, but I wasn't able to. (I also wasn't able to back up the opposite.) Definitive info would be most welcome! 🙂
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe this is related? https://mypy.readthedocs.io/en/stable/generics.html#generic-type-aliases
Co-authored-by: Andreas Klöckner <inform@tiker.net>
abandon (almost) all sets ye who enter here
Closes #465.
Closes #498.
Things to test:
DirectPredecessorsGetter
necessary? Edit: In my tests, it didn't seem to have made a difference, but from the code structure, it appears thatDirectPredecessorsGetter
needs to be deterministic too; changed from orderedsets to unique tuples in 076a76ePlease squash