Update on "New export API with dynamic shape specifications instead of constraints"

avikchaudhuri · avikchaudhuri · commit f0fd8a1a2f99 · 2023-09-21T15:48:18.000-07:00
Our experience using `constraints` / `dynamic_dim` with the existing export API has found it to be (subjectively) clunky and (objectively) verbose in common cases. This PR implements a new design for the export API that replaces the use of `constraints` / `dynamic_dim` with a new way of specifying dynamic shapes, involving the following concepts: * a constructor `Dim` for first-class named dynamic dimensions with ranges (similar to `functorch.dim`, and analogous to internal symbolic sizes) * a mechanism that uses the above in `export` calls to associate inputs to their dynamic shape specifications (`dynamic_shapes`) Design doc: https://docs.google.com/presentation/d/168U7XK72C_WSsZpGESP6Cho9udh193fi0gfjxCNcJ4E/edit#slide=id.p (Meta-only). Note that we only implement Option 1 in that doc. An older version of this PR also implemented Option 3, which is an alternative way of specifying dynamic shapes using tensor type annotations on the exported callable; but we have moved that to future work for now. See docs for these new features in `torch.export`. The existing `torch.export.export` is modified to use the new API, `torch._export.export__RC__`, whenever `constraints=None`. We have not deprecated the existing API yet, but will do in a follow-up. Constraint violation errors arising through use of the new API will now contain suggested fixes using the new API. No longer do we need to report all specializations for static dimensions and suggest all constraints over dynamic dimensions to fix such errors. Instead, due to the redesign, the suggested fixes are much more concise, only involving modifying the definitions of relevant `Dim`s. Differential Revision: [D48919204](https://our.internmc.facebook.com/intern/diff/D48919204/) cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
diff --git a/docs/source/export.rst b/docs/source/export.rst
@@ -357,6 +357,12 @@ Some additional things to note:
   we see in the equality constraints the tuple specifying that ``arg5_1``
   dimension 0 and ``arg6_1`` dimension 0 are equal.
 
+(An experimental mechanism that is designed to eventually subsume the use of
+:func:`torch.export.dynamic_dim` and ``constraints`` involves constructing
+dynamic shape specifications with the :func:`torch.export.Dim` and
+:func:`torch.export.dims` APIs and passing them into :func:`torch.export.export`
+through the ``dynamic_shapes`` argument.)
+
 
 Serialization
 ^^^^^^^^^^^^^
@@ -555,6 +561,8 @@ API Reference
 .. autofunction:: save
 .. autofunction:: load
 .. autofunction:: register_dataclass
+.. autofunction:: Dim
+.. autofunction:: dims
 .. autoclass:: Constraint
 .. autoclass:: ExportedProgram
 
