mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
b27c3558a4
At high current implementation of constrains functions (constrain_as_**) will raise exception for the following code snippets:
```
def f(x):
a = x.item()
constrain_as_size(a, 4, 7)
return torch.empty((a, 4))
inp = torch.tensor([5])
ep = torch._export.export(f, (inp,))
```
The reason is because current constrain logic is:
1) Purely python so it won't survive AOT export (the full node is gone after AOT export since AOT export only maintains aten level op).
2) Utilize side effect to add range constraints for traced symbol's shape env ([code](9591e52880/torch/fx/experimental/symbolic_shapes.py (L370-L372))).
3) If runtime assertion is turned on (by default). [`_AddRuntimeAssertionsForConstraintsPass`](9591e52880/torch/_export/passes/add_runtime_assertions_for_constraints_pass.py (L98-L100)) will try to append assertion node based on range constrains extracted from shape env of symbol during another interpretation round.
4). However, since 1), in the round of AOT export, range constraints logic won't run for symbols generated during this round. And later there is no range constrains information available for assertion round and caused issue.
5) As a result of above, it will failure at `torch.empty((a, 4))` (there is no constrains for `a` that it must be positive).
The fix here is just to implement range constrain logic as a native aten op (CPU implementation as no-op) to make it be able to survive AOT export.
**NOTE:**
[Logic](2d745b95d7/torch/fx/experimental/symbolic_shapes.py (L350-L365C15)) within [`constrain_range`](2d745b95d7/torch/fx/experimental/symbolic_shapes.py (LL313C74-L313C74)) is split out as `constrain_range_int` to capture case when non `SymInt` is passed in and reused in the new `_constrain_range`. The reason is when non `SymInt` is provided:
* If it directly calls `sym_constrain_range`, the C++ version will be called which will be no-op.
* So in this case it calls `constrain_range_int` instead to be able to capture issue like user provides a input whose tensor's shape could be out of range during exporting, like the following for above code example:
```
...
inp = torch.tensor([10])
ep = torch._export.export(f, (inp,)) # immediately raise error
```
Differential Revision: [D46734204](https://our.internmc.facebook.com/intern/diff/D46734204)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/103346
Approved by: https://github.com/tugsbayasgalan
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
from typing import Optional, Callable, Union
|
|
|
|
import torch
|
|
from torch import SymInt, SymFloat
|
|
from torch._dynamo import allow_in_graph
|
|
from torch.fx.experimental.symbolic_shapes import constrain_range_int
|
|
from torch.utils._sympy.value_ranges import ValueRangeError
|
|
|
|
# `Scalar` type used in native_functions.ymal will be translated to `Union[Number, _complex]`
|
|
# could cause type error during since `SymInt` or `SymFloat` will be used.
|
|
# Here manually specify the type explicitly.
|
|
sym_constrain_range: Callable[
|
|
[Union[int, float, SymInt, SymFloat], Optional[int], Optional[int]],
|
|
None,
|
|
] = torch.sym_constrain_range # type: ignore[assignment]
|
|
|
|
|
|
# TODO: we want to hide this min/max stuff under some abstraction similar to
|
|
# DynamicDim
|
|
@allow_in_graph
|
|
def constrain_as_value(symbol, min: Optional[int] = None, max: Optional[int] = None):
|
|
"""
|
|
Add min/max constraint on the intermediate symbol at tracing time
|
|
"""
|
|
|
|
if not isinstance(symbol, SymInt):
|
|
constrain_range_int(symbol, min=min, max=max)
|
|
else:
|
|
sym_constrain_range(symbol, min, max)
|
|
|
|
return symbol
|
|
|
|
|
|
# TODO: we want to hide this min/max stuff under some abstraction similar to
|
|
# DynamicDim
|
|
@allow_in_graph
|
|
def constrain_as_size(symbol, min: int = 2, max: Optional[int] = None):
|
|
"""
|
|
Add min/max constraint on the intermediate symbol which will be used as a size
|
|
"""
|
|
|
|
# TODO: we should investigate turning off 0/1 specialization for unbacked
|
|
# SymInts
|
|
if min < 2:
|
|
raise ValueRangeError(
|
|
"Unable to set min size to be <= 2 because we specialize on 0/1 sizes."
|
|
)
|
|
return constrain_as_value(symbol, min, max)
|