[inductor] propagate shapes in CSEVariable #152198
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/152198
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 033bfb7 with merge base d119481 ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
| @@ -74,6 +74,8 @@ | |||
| # causes typing errors in subclasses (defined in other files). | |||
| OpVarT = str | |||
|
|
|||
| ShapeType = Optional[Sequence[Union[int, str]]] | |||
There was a problem hiding this comment.
There's already a ShapeType in _prims_common defined as Union[torch.Size, list[int], tuple[int, ...]] and this type is already in use in _inductor.
What does it mean if a Shape has strings? Does it make sense to have str and int in the same collection?
Also, perhaps/probably the base type shouldn't be Optional? I see there are also Optional[ShapeType] type annotations below...?
| @@ -1761,6 +1765,7 @@ def generate( | |||
| write: bool = True, | |||
| assignment: bool = True, | |||
| dtype: Optional[torch.dtype] = None, | |||
| shape: Optional[ShapeType] = None, | |||
There was a problem hiding this comment.
Does it have to be Optional here, either? Couldn't the default be ()?
It makes the code simpler if you're not always checking for None.
There was a problem hiding this comment.
() means the variable is a scalar. None means we don't know.
[ghstack-poisoned]
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov