Add type annotations around boxtree.fmm - #125
Conversation
| global_tgt_idx_all_ranks=None): | ||
| src_weight_vecs: Sequence[Array], *, | ||
| global_src_idx_all_ranks: Sequence[Array] | None = None, | ||
| global_tgt_idx_all_ranks: Sequence[Array] | None = None) -> Array | None: |
There was a problem hiding this comment.
drive_fmm can return None in the distributed case (in workers).
There was a problem hiding this comment.
I agree, but that's clearly cumbersome, typing-wise. Can we make it so that we recognize the distributed-worker case via the type system?
There was a problem hiding this comment.
I'm not sure that will work very well. The distributed stuff works by inserting a mixin into an existing wrangler (from what I can tell?), so it's not easy for the type system to check that..
I added an overload in https://github.com/inducer/boxtree/compare/88c4d64fc59db123db07642b40aade35d89b63a5..8347612c329f3a6cbb67d95b3b133451992165b3 that should work? Not sure how to test it without also typing the distributed parts and some tests though 🤷 at least pyright doesn't seem too upset about it..
There was a problem hiding this comment.
Also, the overload checks "is DistributedMixin and indices is None", which I'm not super sure is true on the worker side, but we can probably enforce it (?).
There was a problem hiding this comment.
I added an overload
That's totally good enough, thanks!
There was a problem hiding this comment.
Also, the overload checks "is DistributedMixin and indices is None", which I'm not super sure is true on the worker side, but we can probably enforce it (?).
Type checker is happy, tests pass---happy to snooze this until it becomes relevant.
inducer
left a comment
inducer
left a comment
There was a problem hiding this comment.
Thanks! Just one minor thing regarding dist-mem and None.
|
Thank you! |
2 participants
This adds a bunch of type annotations around the
ExpansionWranglerInterface.