[export] update dynamic shapes section #3214
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/tutorials/3214
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1e157d2 with merge base 33a52a5 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| @@ -190,7 +190,7 @@ def forward(self, x): | |||
| # about safety, but not all Python code is supported, causing these graph | |||
| # breaks. | |||
| # | |||
| # To address this issue, in PyTorch 2.3, we introduced a new mode of | |||
| # To address this issue, in PyTorch 2.5, we introduced a new mode of | |||
| y: torch.Tensor, # [8, 4] | ||
| z: torch.Tensor, # [32] | ||
| ): | ||
| x0 = x + y # output shape: [8, 4] |
There was a problem hiding this comment.
do you need to say "output shape"? IMO it's better to establish the convention in a sentence before the examples ("we'll mention the shapes of local variables in inline comments" or some such thing).
| return x1, x3 | ||
|
|
||
| ###################################################################### | ||
| # By default, ``torch.export`` produces a static program. One clear consequence of this is that at runtime, |
| # Before we look at the program that's produced, let's understand what specifying ``dynamic_shapes`` entails, | ||
| # and how that interacts with export. For every input dimension where a ``Dim`` object is specified, a symbol is | ||
| # allocated, taking on a range of ``[2, inf]`` (why not ``[0, inf]`` or ``[1, inf]``? we'll explain later in the | ||
| # 0/1 specialization section). |
There was a problem hiding this comment.
IMO it's worth linking to https://pytorch.org/docs/main/export.programming_model.html#basics-of-symbolic-shapes when you talk about "a symbol is allocated."
| ): | ||
| x0 = x + y # guard: s2 == s4 | ||
| x1 = self.l(w) # guard: s1 == 5 | ||
| x2 = x0.flatten() |
There was a problem hiding this comment.
mention no guard added here; also may be worth adding the shapes like above, only now symbolic...
| "x": {0: inp4_dim0, 1: inp4_dim1}, | ||
| "y": {0: inp5_dim0, 1: inp5_dim1}, | ||
| } | ||
| # Here you might ask why export "specializes"; why we resolve this static/dynamic conflict by going with the static route. The answer is because |
| # treating this shape as a concrete integer 4. A symbol is allocated for ``y.shape[1]``, and so we finally emit the guard ``s3 == 4``, leading to | ||
| # specialization. | ||
| # | ||
| # One feature of export is that during tracing, statements like asserts, ``torch._checks()``, and ``if/else`` conditions will also emit guards. |
There was a problem hiding this comment.
_checks() is not a thing, _check() is
| # Additionally, you might ask why we traced only the ``if`` branch, and if it's possible to maintain control-flow in your program and keep both branches | ||
| # alive. For that, refer to rewriting your model code following the ``Control Flow Ops`` section above. | ||
| # | ||
| # Since we're talking about guards and specializations, it's a good time to talk about the 0/1 specialization issue we brought up earlier. |
There was a problem hiding this comment.
Worth structuring this as a subsection or something with an anchor, and referencing (linking) to it when you mentioned it at the beginning?
| ) | ||
| ep.module()(torch.randn(2, 4)) | ||
|
|
||
| ###################################################################### |
There was a problem hiding this comment.
So this is a series of examples of increasing complexity that reads like a novel. Probably good to break it up into subsections with meaningful headers so that readers can pause / resume.
| # | ||
| # - ``None`` is a good option for static behavior: | ||
| # - ``dynamic_shapes=None`` (default) exports with the entire model being static. | ||
| # - specifying ``None`` at an input-level exports with all tensor dimensions static, and alternatively is also required for non-tensor inputs. |
…/tutorials into pianpwk/export_ds_1_2025
Description
Updates the export dynamic shapes section to be more up-to-date with the typical user workflow
Checklist
cc @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4