Replace `constraints` with `dynamic_shapes` in export-to-executorch tutorial (#117916)
Summary: `constraints` argument for `torch.export` has been deprecated in favor of the `dynamic_shapes` argument. This PR updates the use of the deprecated API in export-to-executorch tutorial.
Test Plan: CI
Differential Revision: D52932772
Pull Request resolved: https://github.com/pytorch/pytorch/pull/117916
Approved by: https://github.com/angelayi, https://github.com/avikchaudhuri
diff --git a/torch/_export/__init__.py b/torch/_export/__init__.py
index da38079..8ec293f 100644
--- a/torch/_export/__init__.py
+++ b/torch/_export/__init__.py
@@ -97,6 +97,7 @@
args: Tuple[Any],
kwargs: Optional[Dict[str, Any]] = None,
constraints: Optional[List[Constraint]] = None,
+ dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
) -> torch.nn.Module:
"""
A helper function that is intended to trace a module before any pre-autograd
@@ -111,18 +112,50 @@
kwargs: optional example keyword inputs.
- constraints: A optional list of constraints on the dynamic arguments specifying
- their possible range of their shapes
+ constraints: [DEPRECATED: use ``dynamic_shapes`` instead, see below]
+ An optional list of constraints on the dynamic arguments
+ that specify their possible range of shapes. By default, shapes of
+ input torch.Tensors are assumed to be static. If an input torch.Tensor
+ is expected to have dynamic shapes, please use :func:`dynamic_dim`
+ to define :class:`Constraint` objects that specify the dynamics and the possible
+ range of shapes. See :func:`dynamic_dim` docstring for examples on
+ how to use it.
+
+ dynamic_shapes: Should either be:
+ 1) a dict from argument names of ``f`` to their dynamic shape specifications,
+ 2) a tuple that specifies dynamic shape specifications for each input in original order.
+ If you are specifying dynamism on keyword args, you will need to pass them in the order that
+ is defined in the original function signature.
+
+ The dynamic shape of a tensor argument can be specified as either
+ (1) a dict from dynamic dimension indices to :func:`Dim` types, where it is
+ not required to include static dimension indices in this dict, but when they are,
+ they should be mapped to None; or (2) a tuple / list of :func:`Dim` types or None,
+ where the :func:`Dim` types correspond to dynamic dimensions, and static dimensions
+ are denoted by None. Arguments that are dicts or tuples / lists of tensors are
+ recursively specified by using mappings or sequences of contained specifications.
Returns:
An nn.Module containing the traced method.
"""
from torch.export._trace import _convert_input_to_fake, DEFAULT_EXPORT_DYNAMO_CONFIG
+ from torch.export.dynamic_shapes import _process_dynamic_shapes
if kwargs is None:
kwargs = {}
+ if constraints is not None:
+ warnings.warn(
+ "Using `constraints` to specify dynamic shapes for export is DEPRECATED "
+ "and will not be supported in the future. "
+ "Please use `dynamic_shapes` instead (see docs on `torch.export.export`).",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ else:
+ constraints = _process_dynamic_shapes(f, args, kwargs, dynamic_shapes)
+
decomp_table = {op: op.decompose for op in FunctionalTensor.maybe_aliasing_or_mutating_ops}
with torch._dynamo.config.patch(dataclasses.asdict(DEFAULT_EXPORT_DYNAMO_CONFIG)):
m = torch._dynamo.export(
