fix example

Author: xadupre

_doc/recipes/plot_dynamic_shapes_max.py

@@ -126,11 +126,13 @@ def register(fct, fct_shape, namespace, fname):
 # %%
 # Now everything is registered. Let's rewrite the model.
 
+
 class RewrittenModel(torch.nn.Module):
     def forward(self, x, y, fact):
         z = torch.ops.mylib.copy_max_dimensions(x, y)
         return z * fact
 
+
 # %%
 # And check it works.
 

_doc/recipes/plot_dynamic_shapes_nonzero.py

@@ -2,7 +2,7 @@
 Half certain nonzero
 ====================
 
-:func:`torch.nonzero` returns the indices or the first zero found
+:func:`torch.nonzero` returns the indices of the first zero found
 in a tensor. The output shape is unknown in the generic case
 but... If you have a 2D tensor with at least a nonzero value
 in every row, you can guess the dimension. But :func:`torch.export.export`
@@ -49,7 +49,7 @@ def forward(self, x):
 # ++++++
 
 DYN = torch.export.Dim.DYNAMIC
-ep = torch.export.export(model, (x,), dynamic_shapes=((DYN, DYN),))
+ep = torch.export.export(model, (x,), dynamic_shapes=(({0: DYN, 1: DYN}),))
 print(ep)
 
 

_doc/recipes/plot_dynamic_shapes_python_int.py

@@ -36,7 +36,7 @@ def forward(self, x):
 # ++++++
 
 DYN = torch.export.Dim.DYNAMIC
-ep = torch.export.export(model, (x,), dynamic_shapes=((DYN, DYN),))
+ep = torch.export.export(model, (x,), dynamic_shapes=(({0: DYN, 1: DYN}),))
 print(ep)
 
 # %%
@@ -65,7 +65,7 @@ def forward(self, x):
 # Export
 # ++++++
 
-ep = torch.export.export(rewritten_model, (x,), dynamic_shapes=((DYN, DYN),))
+ep = torch.export.export(rewritten_model, (x,), dynamic_shapes=({0: DYN, 1: DYN},))
 print(ep)
 
 
@@ -79,7 +79,7 @@ def forward(self, x):
 
 
 with bypass_export_some_errors(stop_if_static=True):
-    ep = torch.export.export(model, (x,), dynamic_shapes=((DYN, DYN),))
+    ep = torch.export.export(model, (x,), dynamic_shapes=({0: DYN, 1: DYN},))
     print(ep)
 
 # %%