mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
6a1b820255
With `AdditionalInputs`, the behavior is the same as with tensors:
```python
class M(torch.nn.Module):
def forward(self, x, y):
return x + y
additional_inputs = torch.export.AdditionalInputs()
additional_inputs.add((5, 5))
additional_inputs.add((3, 5))
additional_inputs.add((5, 4))
ep = torch.export.export(
M(), (6, 7), dynamic_shapes=additional_inputs, strict=False
)
```
With `ShapesCollection`, we now need to wrap integer inputs as `_IntWrapper` so that we can have a unique identifier for each integer input.
```python
class M(torch.nn.Module):
def forward(self, x, y):
return x + y
from torch.export.dynamic_shapes import _IntWrapper
args = (_IntWrapper(5), _IntWrapper(5))
# Or we can do `args = pytree.tree_map_only(int, lambda a: _IntWrapper(a), orig_args)`
shapes_collection = torch.export.ShapesCollection()
shapes_collection[args[0]] = Dim.DYNAMIC
shapes_collection[args[1]] = Dim.DYNAMIC
ep = torch.export.export(
M(), args, dynamic_shapes=shapes_collection, strict=False
)
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/151842
Approved by: https://github.com/pianpwk
|
||
|---|---|---|
| .. | ||
| experimental | ||
| passes | ||
| __init__.py | ||
| _draft_export.py | ||
| _remove_auto_functionalized_pass.py | ||
| _remove_effect_tokens_pass.py | ||
| _safeguard.py | ||
| _swap.py | ||
| _trace.py | ||
| _tree_utils.py | ||
| _unlift.py | ||
| _wrapper_utils.py | ||
| custom_obj.py | ||
| custom_ops.py | ||
| decomp_utils.py | ||
| dynamic_shapes.py | ||
| exported_program.py | ||
| graph_signature.py | ||
| unflatten.py | ||