Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion _doc/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down
2 changes: 2 additions & 0 deletions _doc/recipes/plot_dynamic_shapes_nonzero.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
"""
.. _l-plot-nonzero:

Half certain nonzero
====================

Expand Down
64 changes: 64 additions & 0 deletions _doc/recipes/plot_export_dim1.py
Original file line number Diff line number Diff line change
@@ -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")
Loading