diff --git a/_doc/index.rst b/_doc/index.rst index 95a3f025..88dbabf7 100644 --- a/_doc/index.rst +++ b/_doc/index.rst @@ -71,11 +71,17 @@ Enlightening Examples * :ref:`l-plot-export_tiny_phi2` -**Torch Export** +**Exporter Recipes** +* :ref:`l-plot-export-dim1` * :ref:`l-plot-export-cond` * :ref:`l-plot-export-with-dynamic` * :ref:`l-plot-export-with-dynamic-shape` +* :ref:`l-plot-dynamic-shapes-python-int` + +**Torch Export Models** + +* :ref:`l-plot-nonzero` * :ref:`l-plot-export-locale-issue` * :ref:`l-plot-tiny-llm-export` * :ref:`l-plot-tiny-llm-export-patched` diff --git a/_doc/recipes/plot_dynamic_shapes_nonzero.py b/_doc/recipes/plot_dynamic_shapes_nonzero.py index 0e6fdbf7..5a8f0e6e 100644 --- a/_doc/recipes/plot_dynamic_shapes_nonzero.py +++ b/_doc/recipes/plot_dynamic_shapes_nonzero.py @@ -1,4 +1,6 @@ """ +.. _l-plot-nonzero: + Half certain nonzero ==================== diff --git a/_doc/recipes/plot_export_dim1.py b/_doc/recipes/plot_export_dim1.py new file mode 100644 index 00000000..20a85abb --- /dev/null +++ b/_doc/recipes/plot_export_dim1.py @@ -0,0 +1,64 @@ +""" +.. _l-plot-export-dim1: + +0, 1, 2 for a Dynamic Dimension in the dummy example to export a model +====================================================================== + +:func:`torch.export.export` does not work if a tensor given to the function +has 0 or 1 for dimension declared as dynamic dimension. + +Simple model, no dimension with 0 or 1 +++++++++++++++++++++++++++++++++++++++ +""" + +import torch +from onnx_diagnostic import doc + + +class Model(torch.nn.Module): + def forward(self, x, y, z): + return torch.cat((x, y), axis=1) + z + + +model = Model() +x = torch.randn(2, 3) +y = torch.randn(2, 5) +z = torch.randn(2, 8) +model(x, y, z) + +DYN = torch.export.Dim.DYNAMIC +ds = {0: DYN, 1: DYN} + +ep = torch.export.export(model, (x, y, z), dynamic_shapes=(ds, ds, ds)) +print(ep.graph) + +# %% +# Same model, a dynamic dimension = 1 +# +++++++++++++++++++++++++++++++++++ + +z = z[:1] + +DYN = torch.export.Dim.DYNAMIC +ds = {0: DYN, 1: DYN} + +try: + ep = torch.export.export(model, (x, y, z), dynamic_shapes=(ds, ds, ds)) + print(ep.graph) +except Exception as e: + print("ERROR", e) + +# %% +# It failed. Let's try a little trick. + +# %% +# Same model, a dynamic dimension = 1 and backed_size_oblivious=True +# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + +with torch.fx.experimental._config.patch(backed_size_oblivious=True): + ep = torch.export.export(model, (x, y, z), dynamic_shapes=(ds, ds, ds)) + print(ep.graph) + +# %% +# It worked. + +doc.plot_legend("dynamic dimension\nworking with\n0 or 1", "torch.export.export", "green")