[FX] Fix PyTree unpacking carrying forward type annotations #81906
Conversation
[ghstack-poisoned]
🔗 Helpful links
❌ 1 New FailuresAs of commit 634af35 (more details on the Dr. CI page): Expand to see more
🕵️ 1 new failure recognized by patternsThe following CI failures do not appear to be due to upstream breakages
|
Resolves #81902 [ghstack-poisoned]
Resolves #81902 [ghstack-poisoned]
There was a problem hiding this comment.
LGTM.
But I wonder whether we should rethink our pytree integration with FX. Like, our original integration had (essentially) 2 requirements.
- The FX module was still runnable as-is with the same inputs as before.
- The FX module is retraceable with
fx.symbolic_trace.
From my understanding, most of the users of pytrees are now going through make_fx - can we drop some of the previous requirements and simplify this code?
| function_definition = super().gen_fn_def(function_args[:], maybe_return_annotation) | ||
| if len(free_vars) > 0: # pytree has placeholders in it | ||
| function_definition += f""" | ||
| {', '.join(free_vars)}, = fx_pytree.tree_flatten_spec([{', '.join(function_args)}], self._in_spec)""" | ||
| {', '.join(var.name for var in free_vars)}, = fx_pytree.tree_flatten_spec\ |
There was a problem hiding this comment.
why trailing comma before = ?
There was a problem hiding this comment.
That forces it to always be a tuple even if there's 1 element
|
@pytorchbot merge |
|
@pytorchbot successfully started a merge job. Check the current status here |
|
Merge failed due to Refusing to merge as mandatory check(s) pull failed for rule superuser |
|
@pytorchbot merge -f |
|
@pytorchbot merge -f "flaky test" |
|
@pytorchbot successfully started a merge job. Check the current status here |
|
Hey @jamesr66a. |
|
@pytorchbot revert -m "breaking internal builds" -c "ghfirst" |
|
@pytorchbot successfully started a revert job. Check the current status here |
|
@jamesr66a your PR has been successfully reverted. |
…81906)" This reverts commit e0d83a0. Reverted #81906 on behalf of https://github.com/jeanschmidt due to breaking internal builds
|
@jamesr66a I saw this got reverted, but we still need this fix. @jeanschmidt how is it breaking internal builds? |
Stack from ghstack:
Resolves #81902