File size: 54,785 Bytes
f4cade0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207 1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 |
# mypy: allow-untyped-defs
import dataclasses
import inspect
import logging
import sys
from collections import defaultdict
from enum import auto, Enum
from typing import Any, Callable, Optional, TYPE_CHECKING, Union
import torch
from torch.utils._pytree import (
_get_node_type,
BUILTIN_TYPES,
keystr,
LeafSpec,
MappingKey,
SequenceKey,
SUPPORTED_NODES,
tree_flatten,
tree_map,
tree_map_with_path,
)
from .exported_program import ExportedProgram
if TYPE_CHECKING:
from sympy import Symbol
from torch._guards import Source
from torch.fx.experimental.symbolic_shapes import ShapeEnv, StrictMinMaxConstraint
__all__ = [
"Constraint",
"Dim",
"dims",
"refine_dynamic_shapes_from_suggested_fixes",
"AdditionalInputs",
]
log = logging.getLogger(__name__)
class _DimHintType(Enum):
"""
Enum for dynamic shape hints.
- AUTO means automatic inference of shape (static or dynamic).
- STATIC means static shape (always specialized).
- DYNAMIC means dynamic, will error out if specialized.
"""
AUTO = auto()
STATIC = auto()
DYNAMIC = auto()
@dataclasses.dataclass
class _DimHint:
type: _DimHintType
min: Optional[int] = None
max: Optional[int] = None
_factory: Optional[bool] = True
@staticmethod
def AUTO():
return _DimHint(_DimHintType.AUTO)
@staticmethod
def DYNAMIC():
return _DimHint(_DimHintType.DYNAMIC)
@staticmethod
def STATIC():
return _DimHint(_DimHintType.STATIC)
def __call__(self, min=None, max=None) -> "_DimHint":
if not self._factory:
raise TypeError(f"'{type(self)}' object is not callable")
assert min is None or min >= 0, "min must be non-negative"
assert max is None or max >= 0, "max must be non-negative"
assert min is None or max is None or min <= max, "min must be <= max"
return _DimHint(self.type, min=min, max=max, _factory=False)
class Dim:
"""
The `Dim` class allows users to specify dynamism in their exported programs. By marking a dimension with a `Dim`,
the compiler associates the dimension with a symbolic integer containing a dynamic range.
The API can be used in 2 ways: Dim hints (i.e. automatic dynamic shapes: `Dim.AUTO`, `Dim.DYNAMIC`, `Dim.STATIC`),
or named Dims (i.e. `Dim("name", min=1, max=2)`).
Dim hints provide the lowest barrier to exportability, with the user only needing to specify if a dimension
if dynamic, static, or left for the compiler to decide (`Dim.AUTO`). The export process will automatically
infer the remaining constraints on min/max ranges and relationships between dimensions.
Example::
class Foo(nn.Module):
def forward(self, x, y):
assert x.shape[0] == 4
assert y.shape[0] >= 16
return x @ y
x = torch.randn(4, 8)
y = torch.randn(8, 16)
dynamic_shapes = {
"x": {0: Dim.AUTO, 1: Dim.AUTO},
"y": {0: Dim.AUTO, 1: Dim.AUTO},
}
ep = torch.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)
Here, export would raise an exception if we replaced all uses of `Dim.AUTO` with `Dim.DYNAMIC`,
as x.shape[0] is constrained to be static by the model.
More complex relations between dimensions may also be codegened as runtime assertion nodes by the compiler,
e.g. (x.shape[0] + y.shape[1]) % 4 == 0, to be raised if runtime inputs do not satisfy such constraints.
You may also specify min-max bounds for Dim hints, e.g. `Dim.AUTO(min=16, max=32)`, `Dim.DYNAMIC(max=64)`,
with the compiler inferring the remaining constraints within the ranges. An exception will be raised if
the valid range is entirely outside the user-specified range.
Named Dims provide a stricter way of specifying dynamism, where exceptions are raised if the compiler
infers constraints that do not match the user specification. For example, exporting the previous
model, the user would need the following `dynamic_shapes` argument::
s0 = Dim("s0")
s1 = Dim("s1", min=16)
dynamic_shapes = {
"x": {0: 4, 1: s0},
"y": {0: s0, 1: s1},
}
ep = torch.export(Foo(), (x, y), dynamic_shapes=dynamic_shapes)
Named Dims also allow specification of relationships between dimensions, up to univariate linear relations.
For example, the following indicates one dimension is a multiple of another plus 4::
s0 = Dim("s0")
s1 = 3 * s0 + 4
"""
AUTO = _DimHint.AUTO()
DYNAMIC = _DimHint.DYNAMIC()
STATIC = _DimHint.STATIC()
def __init__(
self, name: str, *, min: Optional[int] = None, max: Optional[int] = None
):
from torch.utils._sympy.numbers import int_oo
_min = 0 if min is None else min
_max = int_oo if max is None else max
assert _max > _min, f"Cannot create Dim with inconsistent min={min}, max={max}"
assert name.isidentifier(), f"Dim name must be a valid identifier, got {name}"
self.__name__ = name
self.min = _min
self.max = _max
def __add__(self, other) -> "Dim":
# e.g., dim + 1
if type(other) is not int:
raise NotImplementedError(
f"Attempted to add {other} to {self.__name__}, where an integer was expected. "
"(Only increasing linear operations with integer coefficients are supported.)"
)
return self._derive(lambda x: x + other)
def __radd__(self, other) -> "Dim":
return self + other
def __sub__(self, other) -> "Dim":
# e.g., dim - 1
if type(other) is not int:
raise NotImplementedError(
f"Attempted to subtract {other} from {self.__name__}, where an integer was expected. "
"(Only increasing linear operations with integer coefficients are supported.)"
)
return self._derive(lambda x: x - other)
def __rsub__(self, other) -> "Dim":
raise NotImplementedError(
f"Attempted to negate {self.__name__}. "
"(Only increasing linear operations with integer coefficients are supported.)"
)
def __mul__(self, other) -> "Dim":
# e.g., dim * 2
if type(other) is not int or other <= 0:
raise NotImplementedError(
f"Attempted to multiply {other} with {self.__name__}, where a positive integer was expected. "
"(Only increasing linear operations with integer coefficients are supported.)"
)
return self._derive(lambda x: x * other)
def __rmul__(self, other) -> "Dim":
return self * other
def _derived_name(self, fn) -> str:
from sympy import sympify
return str(fn(sympify(self.__name__)))
def _derive(self, fn) -> "Dim":
return _DerivedDim(self._derived_name(fn), self, fn)
@staticmethod
def _readable(name: str, min_: int, max_: int) -> str:
from torch.utils._sympy.numbers import int_oo
if min_ == 2:
min_ = None # type: ignore[assignment]
if max_ == int_oo:
max_ = None # type: ignore[assignment]
if min_ is None and max_ is None:
return f"Dim('{name}')"
if min_ is None:
return f"Dim('{name}', max={max_})"
if max_ is None:
return f"Dim('{name}', min={min_})"
return f"Dim('{name}', min={min_}, max={max_})"
def __repr__(self):
return Dim._readable(self.__name__, self.min, self.max)
_Dim = Dim # TODO(pianpwk): remove after it's no longer internally breaking
class _StaticDim(Dim):
"""
Class for static :func:`Dim` types.
This class is only for setting and checking static dim constraints,
and the user should never interact with it.
"""
def __init__(self, value: int):
self.__name__ = str(value)
self.value = value
@property
def min(self): # type: ignore[override]
return self.value # type: ignore[attr-defined]
@property
def max(self): # type: ignore[override]
return self.value # type: ignore[attr-defined]
class _DerivedDim(Dim):
"""
Class for derived :func:`Dim` types.
Currently we only support increasing linear expressions with integer coefficients.
In other words, a derived Dim can always be written in the form Ax + B, where
x is a regular Dim (i.e., non-derived Dim), A and B are integers, and A is positive.
(In particular, the latter ensures that x < y => Ax + B < Ay + B.)
These restrictions on the form of derived Dims makes the metatheory simpler: e.g.,
it simplifies computing ranges for derived Dims, solving for underlying regular Dims,
deciding equalities between derived Dims, and so on.
The function lambda x: Ax + B is expressed by `fn`, where x is a normal Dim, `root`.
The range of a derived Dim is computed by mapping `fn` over the range of its `root`.
"""
def __init__(self, name: str, root: Dim, fn: Callable):
self.__name__ = name
self.root = root
self.fn = fn
@property
def min(self): # type: ignore[override]
# assume that self.fn is an increasing function
# TODO(avik): use sympy value range analysis instead?
from sympy import Integer
from torch.utils._sympy.numbers import int_oo
if self.root.min is -int_oo: # type: ignore[attr-defined]
return -int_oo # fn not needed cuz increasing
_min_symint = self.fn(Integer(self.root.min)) # type: ignore[attr-defined]
root = self.root # type: ignore[attr-defined]
assert _min_symint >= 0, (
f"Expected derived min value of {self.__name__} to be >= 0. "
f"Please specify an appropriate min value for {root.__name__} "
f"(currently {root.min})."
)
return int(_min_symint)
@property
def max(self): # type: ignore[override]
# assume that self.fn is an increasing function
# TODO(avik): use sympy value range analysis instead?
from sympy import Integer
from torch.utils._sympy.numbers import int_oo
if self.root.max is int_oo: # type: ignore[attr-defined]
return int_oo # fn not needed cuz increasing
_max_symint = self.fn(Integer(self.root.max)) # type: ignore[attr-defined]
root = self.root # type: ignore[attr-defined]
assert _max_symint <= sys.maxsize - 1, (
f"Expected derived max value of {self.__name__} to be <= {sys.maxsize - 1}. "
f"Please specify an appropriate max value for {root.__name__} "
f"(currently {root.max})."
)
return int(_max_symint)
def _derive(self, fn):
# We support nesting, e.g., 2*dim + 1.
# This is implemented by composing operations on the same root.
# As a consequence, roots are always regular Dims (i.e., not derived Dims).
return _DerivedDim(
self._derived_name(fn),
self.root,
lambda x: fn(self.fn(x)),
)
def __repr__(self):
return self.__name__
def dims(
*names: str, min: Optional[int] = None, max: Optional[int] = None
) -> tuple[Dim, ...]:
"""
Util to create multiple :func:`Dim` types.
Returns:
A tuple of :func:`Dim` types.
"""
return tuple(Dim(name, min=min, max=max) for name in names) # type: ignore[misc]
@dataclasses.dataclass
class _ConstraintTarget:
"""
This represents input tensor dimensions.
"""
t_id: int
dim: int
@dataclasses.dataclass
class _Constraint(_ConstraintTarget):
"""
This represents a Dim describing a constraint target.
`name` is the name of the Dim.
`constraint_range` contains the min/max bounds of the Dim.
"""
name: str
constraint_range: "StrictMinMaxConstraint"
def _clone_with_range(self, lower=0, upper=None):
# Import sympy locally
from torch.fx.experimental.symbolic_shapes import StrictMinMaxConstraint
from torch.utils._sympy.numbers import int_oo
from torch.utils._sympy.value_ranges import ValueRanges
if upper is None:
upper = int_oo
constraint_range = StrictMinMaxConstraint(
vr=self.constraint_range.vr & ValueRanges(lower=lower, upper=upper),
warn_only=False,
)
return _Constraint(
self.t_id,
self.dim,
self.name,
constraint_range,
)
def __ge__(self, lower):
return self._clone_with_range(lower=lower)
def __gt__(self, lower):
return self._clone_with_range(lower=lower + 1)
def __le__(self, upper):
return self._clone_with_range(upper=upper)
def __lt__(self, upper):
return self._clone_with_range(upper=upper - 1)
def __bool__(self):
# NOTE(avik): We do not support compound expressions like a <= x <= b.
# This is because Python implicitly desugars them into bool(a <= x) and bool(x <= b),
# and moreover, enforces that any overload of __bool__ must return True or False.
# FWIW, sympy also raises TypeError in this case.
raise TypeError(
"Cannot determine truth value of _Constraint. "
"If you are trying to combine _Constraint's with logical connectives, "
"you can specify them separately instead."
)
@property
def serializable_spec(self):
# We need a serialization compatible format of the constraint so that it
# can be savedin the graph module w/o breaking the module serialization.
# The saved constraints will be used directly for the post-exporting pass
# that converts constraints to runtime assertion. The saved constraints
# will not be saved in the serialized module.
# TODO: A better way is needed. Currently we use 't_id' to map the constraint,
# which is not reliable
return {
"t_id": self.t_id,
"dim": self.dim,
"min": self.constraint_range.vr.lower,
"max": self.constraint_range.vr.upper,
}
@dataclasses.dataclass
class _PhantomRoot:
"""
This represents the root of a derived Dim where the root does not directly
specify the shape of any input dimension, but the derived Dim does.
e.g., the input shapes 2*dim and dim + 1 are related via a "phantom" dim.
The fields `name`, `constraint_range`, and `val` carried by a phantom root
help create a symbol for it. Any derived dims with this phantom root are
backed by expressions over this symbol.
"""
name: str
constraint_range: "StrictMinMaxConstraint"
val: int
@dataclasses.dataclass
class _DerivedConstraint(_ConstraintTarget):
"""
This represents a derived Dim, whose root is either a regular constraint target
(which directly specifies the shape of some input dimension) or a phantom root
(which does so indirectly).
It can be thought of as a subclass of `_Constraint`, except that it does not
support <, <=, >, >= operations.
"""
name: str
constraint_range: "StrictMinMaxConstraint"
root: Union[_ConstraintTarget, _PhantomRoot]
fn: Callable
@property
def serializable_spec(self):
# same as _Constraint.serializable_spec
return {
"t_id": self.t_id,
"dim": self.dim,
"min": self.constraint_range.vr.lower,
"max": self.constraint_range.vr.upper,
}
@dataclasses.dataclass
class _RelaxedConstraint(_ConstraintTarget):
"""
This represents a dim marked with Dim.AUTO/DYNAMIC (i.e. mark_dynamic() or maybe_mark_dynamic()),
which leaves relations & min/max ranges for inference, instead of requiring explicit specification.
The intention is for constraint violations to not be raised if produce_guards() finds equalities or
relations between a _RelaxedConstraint and another type of _Constraint.
"""
@property
def serializable_spec(self):
return {
"t_id": self.t_id,
"dim": self.dim,
}
Constraint = Union[_Constraint, _DerivedConstraint, _RelaxedConstraint]
@dataclasses.dataclass
class _IntWrapper:
"""
Dummy wrapper class to wrap around integer inputs so that when we parse the
dynamic_shapes structure, we can mark if any of the integers were marked as
dynamic.
"""
val: int
# Disallow specifying dynamism
dynamism: Optional[Union[_DimHint, int]] = dataclasses.field(
init=False, default=None
)
def _process_equalities(
constraint: Constraint,
get_sources: Callable[[int, int], list["Source"]],
shape_env: "ShapeEnv",
names: dict[str, tuple[int, int]],
source_pairs: list[tuple["Source", "Source"]],
derived_equalities: list[tuple["Source", Union["Source", "Symbol"], Callable]],
phantom_symbols: dict[str, "Symbol"],
relaxed_sources: set["Source"],
):
"""
Updates `source_pairs`, `derived_equalities`, and `phantom_symbols` (which become
fields of `EqualityConstraint`) based on a given input `constraint`.
"""
sources = get_sources(constraint.t_id, constraint.dim)
if not sources: # empty sources due to unused shapes
return
source, *other_sources = sources
# When t.size()[dim] maps to src0, src1, ..., srcN, we add
# constraints that make src0 "equal" to src1, ..., srcN.
source_pairs.extend((source, other_source) for other_source in other_sources)
if isinstance(constraint, _Constraint):
if constraint.name in names:
shared_t_id, shared_dim = names[constraint.name]
other_sources = get_sources(shared_t_id, shared_dim)
source_pairs.extend(
(source, other_source) for other_source in other_sources
)
else:
names[constraint.name] = (constraint.t_id, constraint.dim)
elif isinstance(constraint, _DerivedConstraint):
# branch based on the root of the _DerivedConstraint
if not isinstance(constraint.root, _PhantomRoot):
# either root points to an input source
root = get_sources(constraint.root.t_id, constraint.root.dim)[0]
else:
# or root points to a phantom symbol
if constraint.root.name in phantom_symbols:
root = phantom_symbols[constraint.root.name]
else:
# create a phantom symbol in the shape env based on the _PhantomRoot
root = shape_env.create_symbol(
val=constraint.root.val,
source=torch._dynamo.source.ConstantSource(constraint.root.name),
dynamic_dim=torch.fx.experimental.symbolic_shapes.DimDynamic.DYNAMIC,
constraint_dim=constraint.root.constraint_range,
)
phantom_symbols[constraint.root.name] = root
fn = constraint.fn
# A derived equality (source, root, fn) informally corresponds to source = fn(root).
# Here source describes an input and root might describe another input or a phantom symbol.
derived_equalities.append((source, root, fn))
elif isinstance(constraint, _RelaxedConstraint):
relaxed_sources.add(source)
def _tree_map_with_path(
func: Callable[..., Any],
tree: Any,
*dynamic_shapes: Any,
tree_name: Optional[str] = None,
) -> Any:
"""
Customized tree_map for mapping pytrees to dynamic_shapes.
For built-in types (e.g., standard collections) this behaves exactly like tree_map.
OTOH for a user-defined class C registered with pytree, we cannot assume that a C
containing tensors can be mapped to a C containing dynamic shapes (i.e., C may not
be a polymorphic container). In that case we use the flattened form of C instead.
Thus a C(**tensors) that flattens to (**tensors) will map to (**dynamic_shapes).
Args:
func: function to apply to each (int, float, str, bool, None, torch.Tensor)
tree: input pytree
dynamic_shapes: zero or more (typically one) dynamic_shapes to match
Returns:
output pytree mapping func to each (int, float, str, bool, None, torch.Tensor)
"""
def is_leaf(t):
# BUILTIN_TYPES is a subset of SUPPORTED_NODES, the latter being all types
# registered with pytree. Types *not* in BUILTIN_TYPES include primitive types
# (int, float, str, bool, None, torch.Tensor), which are not in SUPPORTED_NODES,
# as well as user-defined classes registered with pytree, which are.
return _get_node_type(t) not in BUILTIN_TYPES
def f(path, t, *dynamic_shapes):
typ = _get_node_type(t)
# typ is not in BUILTIN_TYPES
if typ in SUPPORTED_NODES:
# thus typ is a user-defined class registered with pytree,
# in which case flatten and recurse
return tree_map_with_path(
f,
SUPPORTED_NODES[typ].flatten_fn(t)[0],
*dynamic_shapes,
is_leaf=is_leaf,
)
else:
return func(path, t, *dynamic_shapes)
try:
return tree_map_with_path(f, tree, *dynamic_shapes, is_leaf=is_leaf)
except ValueError as e:
if "mismatch" in e.args[0]:
# When PyTree finds a structural mismatch between tree and dynamic_shapes,
# the error message is unfortunately quite horrible. Let's fix that.
assert dynamic_shapes, "Cannot be a mismatch if there is no dynamic_shapes"
assert tree_name, "Must provide a tree_name when there might be a mismatch"
def _key(type_, context, i):
# derive a PyTree key given the type, context, and child # of a TreeSpec
if type_ is dict:
return MappingKey(context[i])
if type_ in (list, tuple):
assert context is None
return SequenceKey(i)
raise AssertionError(f"Did not expect type {type_}")
def raise_mismatch_error(msg):
from torch._dynamo.exc import UserError, UserErrorType
raise UserError(
UserErrorType.INVALID_INPUT,
f"Detected mismatch between the structure of `{tree_name}` and `dynamic_shapes`: {msg}",
case_name="dynamic_shapes_validation",
)
def _compare(tree, dynamic_shapes, path):
# raise an error at the point where tree and dynamic_shapes differ,
# including the path to that point and the reason for the difference
rendered_path = keystr(path)
if isinstance(tree, LeafSpec):
return
if isinstance(dynamic_shapes, LeafSpec):
raise_mismatch_error(
f"`{tree_name}{rendered_path}` is a {tree.type}, "
f"but `dynamic_shapes{rendered_path}` is not"
)
if tree.type != dynamic_shapes.type:
raise_mismatch_error(
f"`{tree_name}{rendered_path}` is a {tree.type}, "
f"but `dynamic_shapes{rendered_path}` is a {dynamic_shapes.type}"
)
if len(tree.children_specs) != len(dynamic_shapes.children_specs):
raise_mismatch_error(
f"`{tree_name}{rendered_path}` has {len(tree.children_specs)} elements, "
f"but `dynamic_shapes{rendered_path}` has {len(dynamic_shapes.children_specs)} elements"
)
if tree.type is dict:
# context, children could be out of order
if sorted(tree.context) != sorted(dynamic_shapes.context):
raise_mismatch_error(
f"`{tree_name}{rendered_path}` has keys {tree.context}, "
f"but `dynamic_shapes{rendered_path}` has keys {dynamic_shapes.context}"
)
_remap = dict(
zip(dynamic_shapes.context, dynamic_shapes.children_specs)
)
dynamic_shapes_children_specs = [_remap[k] for k in tree.context]
else:
dynamic_shapes_children_specs = dynamic_shapes.children_specs
for i, (tree_, dynamic_shapes_) in enumerate(
zip(tree.children_specs, dynamic_shapes_children_specs)
):
_compare(
tree_,
dynamic_shapes_,
path + [_key(tree.type, tree.context, i)],
)
_, tree_spec = tree_flatten(tree, is_leaf=is_leaf)
for other_tree in dynamic_shapes:
_, other_tree_spec = tree_flatten(other_tree, is_leaf)
_compare(tree_spec, other_tree_spec, [])
raise
def _combine_args(f, args, kwargs, _is_torch_jit_trace=False) -> dict[str, Any]:
# combine args and kwargs following the signature of f, as it happens
# in the body of f when called with *args, **kwargs
if isinstance(f, ExportedProgram):
f = f.module()
if not _is_torch_jit_trace:
signature = (
inspect.signature(f.forward)
if isinstance(f, torch.nn.Module)
else inspect.signature(f)
)
kwargs = kwargs if kwargs is not None else {}
return signature.bind(*args, **kwargs).arguments
return args
class ShapesCollection:
"""
Builder for dynamic_shapes.
Used to assign dynamic shape specifications to tensors that appear in inputs.
This is useful particularly when :func:`args` is a nested input structure, and it's
easier to index the input tensors, than to replicate the structure of :func:`args` in
the :func:`dynamic_shapes` specification.
Example::
args = {"x": tensor_x, "others": [tensor_y, tensor_z]}
dim = torch.export.Dim(...)
dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[tensor_y] = {0: dim * 2}
# This is equivalent to the following (now auto-generated):
# dynamic_shapes = {"x": (dim, dim + 1, 8), "others": [{0: dim * 2}, None]}
torch.export(..., args, dynamic_shapes=dynamic_shapes)
To specify dynamism for integers, we need to first wrap the integers using
_IntWrapper so that we have a "unique identification tag" for each integer.
Example::
args = {"x": tensor_x, "others": [int_x, int_y]}
# Wrap all ints with _IntWrapper
mapped_args = pytree.tree_map_only(int, lambda a: _IntWrapper(a), args)
dynamic_shapes = torch.export.ShapesCollection()
dynamic_shapes[tensor_x] = (dim, dim + 1, 8)
dynamic_shapes[mapped_args["others"][0]] = Dim.DYNAMIC
# This is equivalent to the following (now auto-generated):
# dynamic_shapes = {"x": (dim, dim + 1, 8), "others": [Dim.DYNAMIC, None]}
torch.export(..., args, dynamic_shapes=dynamic_shapes)
"""
def __init__(self):
self._shapes = {}
def __setitem__(self, t, shape):
assert isinstance(t, (torch.Tensor, _IntWrapper)), (
f"Cannot assign shape to non-tensor or non-_IntWrapper type {type(t)}"
)
# TODO(avik): check that shape is indeed a Shape
t_id = id(t)
if t_id in self._shapes:
_shape = self._shapes[t_id]
assert shape == _shape, (
f"Shapes assigned to input do not match: expected {_shape}, got {shape}"
)
else:
self._shapes[id(t)] = shape
def __getitem__(self, t):
t_id = id(t)
if t_id not in self._shapes:
self._shapes[t_id] = {}
return self._shapes[t_id]
def __len__(self):
return len(self._shapes)
def dynamic_shapes(self, m, args, kwargs=None):
"""
Generates the :func:`dynamic_shapes` pytree structure according to :func:`args` and :func:`kwargs`.
"""
t_ids = set()
def find_shape(path, t):
t_id = id(t)
if t_id in self._shapes:
t_ids.add(t_id)
return self._shapes[t_id]
else:
return None
combined_args = _combine_args(m, args, kwargs)
dynamic_shapes = _tree_map_with_path(find_shape, combined_args)
if any(t_id not in t_ids for t_id in self._shapes):
raise ValueError(
"Some tensors that were assigned shapes were not found in args. "
"Maybe such tensors were copied when passing them as args? "
"Maybe such tensors are contained in classes that were not registered with pytree?"
)
return dynamic_shapes
class AdditionalInputs:
"""
Infers dynamic_shapes based on additional inputs.
This is useful particularly for deployment engineers who, on the one hand, may
have access to ample testing or profiling data that can provide a fair sense of
representative inputs for a model, but on the other hand, may not know enough
about the model to guess which input shapes should be dynamic.
Input shapes that are different than the original are considered dynamic; conversely,
those that are the same as the original are considered static. Moreover, we verify
that the additional inputs are valid for the exported program. This guarantees that
tracing with them instead of the original would have generated the same graph.
Example::
args0, kwargs0 = ... # example inputs for export
# other representative inputs that the exported program will run on
dynamic_shapes = torch.export.AdditionalInputs()
dynamic_shapes.add(args1, kwargs1)
...
dynamic_shapes.add(argsN, kwargsN)
torch.export(..., args0, kwargs0, dynamic_shapes=dynamic_shapes)
"""
def __init__(self):
self._examples = []
def add(self, args, kwargs=None):
"""
Additional input :func:`args` and :func:`kwargs`.
"""
assert type(args) is tuple, f"Representative args {args} must be a tuple"
assert kwargs is None or type(kwargs) is dict, (
f"Representative kwargs {kwargs} must be None or a dict"
)
self._examples.append((args, kwargs))
def dynamic_shapes(self, m, args, kwargs=None):
"""
Infers a :func:`dynamic_shapes` pytree structure by merging shapes of the
original input :func:`args` and :func:`kwargs` and of each additional input
args and kwargs.
"""
dynamic_shapes, *other_dynamic_shapes = [
_tree_map_with_path(
lambda path, t: tuple(t.shape) if isinstance(t, torch.Tensor) else t,
_combine_args(m, args, kwargs),
)
for args, kwargs in [(args, kwargs), *self._examples]
]
def _mark_dynamism(v, *other_vs):
if not all(type(v) == type(other) for other in other_vs):
raise ValueError(
"The following inputs were found to have differing types, "
f"so they cannot be marked as dynamic: {(v,) + other_vs}."
)
if isinstance(v, int) and not isinstance(v, bool):
if all(other_v == v for other_v in other_vs):
return None
else:
return Dim.DYNAMIC
else:
if not all(other_v == v for other_v in other_vs):
raise ValueError(
"The following inputs were found to have differing values, "
f"but they cannot be marked as dynamic: {(v,) + other_vs}."
)
return None
return tree_map(
_mark_dynamism,
dynamic_shapes,
*other_dynamic_shapes,
is_leaf=lambda i: type(i) is int,
)
def verify(self, ep):
"""
Verifies that an exported program is valid for each additional input.
"""
epm = ep.module()
for args, kwargs in self._examples:
torch.export._unlift._check_input_constraints_pre_hook(
epm, args, kwargs or {}
)
def _warn_on_None_dynamic_shape_dimension():
msg = (
"Using None as a dynamic shape dimension is deprecated. "
"Please use Dim.STATIC instead"
)
# TODO(avik): raise an error in the future
log.warning(msg)
def _check_dynamic_shapes(
combined_args: dict[str, Any],
dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
):
"""
Checks the dynamic_shapes specification for correctness,
using combined args + kwargs as reference for inputs structure.
"""
from torch._dynamo.exc import UserError, UserErrorType
if dynamic_shapes is None or len(dynamic_shapes) == 0:
return
if isinstance(dynamic_shapes, (tuple, list)):
combined_args = type(dynamic_shapes)(combined_args.values()) # type: ignore[assignment, misc]
bounds: dict[str, tuple[int, int]] = {}
def check_same_bounds(dim):
if dim.__name__ in bounds:
min_, max_ = bounds[dim.__name__]
if dim.min != min_ or dim.max != max_:
this_ = Dim._readable(dim.__name__, min_, max_)
that_ = Dim._readable(dim.__name__, dim.min, dim.max)
raise UserError(
UserErrorType.INVALID_INPUT,
f"Found different definitions {this_} and {that_} "
f"for the same symbolic dimension {dim}!",
)
else:
bounds[dim.__name__] = (dim.min, dim.max)
def check_symbols(path, tensor, shape):
if isinstance(shape, dict):
for i, dim in shape.items():
if isinstance(dim, Dim):
check_same_bounds(dim)
elif dim is None:
_warn_on_None_dynamic_shape_dimension()
elif not (isinstance(dim, (int, _DimHint))):
raise UserError(
UserErrorType.INVALID_INPUT,
f"Unexpected dimension mapped to index {i} in input tensor shape {shape} "
f"specified at `dynamic_shapes{keystr(path)}` "
f"(expected None, an int, a Dim, Dim.AUTO, Dim.STATIC, or Dim.DYNAMIC, "
f" but got {dim} instead)",
case_name="dynamic_shapes_validation",
)
elif isinstance(shape, (tuple, list)):
if len(shape) != len(tensor.shape):
raise UserError(
UserErrorType.INVALID_INPUT,
f"Expected dynamic shape spec {shape} specified at `dynamic_shapes{keystr(path)}` "
f"to have the same length as the actual tensor shape {tensor.shape} "
f"(expected {len(tensor.shape)}, but got {len(shape)} instead)",
case_name="dynamic_shapes_validation",
)
for i, dim in enumerate(shape):
if isinstance(dim, Dim):
check_same_bounds(dim)
elif dim is None:
_warn_on_None_dynamic_shape_dimension()
elif not (isinstance(dim, (int, _DimHint))):
raise UserError(
UserErrorType.INVALID_INPUT,
f"Unexpected dimension #{i} in input tensor shape {shape} "
f"specified at `dynamic_shapes{keystr(path)}` "
f"(expected None, an int, a Dim, Dim.AUTO, Dim.STATIC, or Dim.DYNAMIC, "
f"but got {dim} instead)",
case_name="dynamic_shapes_validation",
)
elif shape is not None:
raise UserError(
UserErrorType.INVALID_INPUT,
f"Unexpected input tensor shape {shape} specified at `dynamic_shapes{keystr(path)}` "
f"(expected either a list/tuple of dimensions, or a dict mapping indices to dimensions,"
f" where each dimension is an int, a Dim, Dim.AUTO, Dim.STATIC, or Dim.DYNAMIC)",
case_name="dynamic_shapes_validation",
)
assert isinstance(dynamic_shapes, (dict, tuple, list))
if isinstance(dynamic_shapes, dict):
got_keys = list(dynamic_shapes.keys())
expected_arg_names = list(combined_args.keys())
if sorted(got_keys) != sorted(expected_arg_names):
msg = (
f"When `dynamic_shapes` is specified as a dict, its top-level keys "
f"must be the arg names {expected_arg_names} of `inputs`, but "
f"here they are {got_keys}. "
)
if (
len(combined_args) == 1
and expected_arg_names[0] not in got_keys
and isinstance(combined_args[expected_arg_names[0]], dict)
):
msg += (
"Since here `inputs` is a list/tuple enclosing a single dict, "
"maybe you just forgot to enclose `dynamic_shapes` in a list/tuple?"
)
else:
msg += (
"Alternatively, you could also ignore arg names entirely "
"and specify `dynamic_shapes` as a list/tuple matching `inputs`."
)
raise UserError(
UserErrorType.INVALID_INPUT, msg, case_name="dynamic_shapes_validation"
)
def check_shape(path, t, dynamic_shape):
if isinstance(t, torch.Tensor):
check_symbols(path, t, dynamic_shape)
elif isinstance(t, _IntWrapper):
if isinstance(dynamic_shape, _Dim):
raise ValueError(
"Unable to specify input integers as dynamic through named "
"Dims. Please use Dim.AUTO/DYNAMIC instead."
)
assert dynamic_shape is None or isinstance(dynamic_shape, (int, _DimHint))
else:
if dynamic_shape is not None:
rendered_path = keystr(path)
raise UserError(
UserErrorType.INVALID_INPUT,
f"Cannot associate shape {dynamic_shape} specified at `dynamic_shapes{rendered_path}` "
f"to non-tensor type {type(t)} at `inputs{rendered_path}` (expected None)",
case_name="dynamic_shapes_validation",
)
_tree_map_with_path(check_shape, combined_args, dynamic_shapes, tree_name="inputs")
def _process_dynamic_shapes(
combined_args: dict[str, Any],
dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
) -> list[Constraint]:
"""
Reads the dynamic_shapes specification and produces a list of constraints.
"""
from torch._dynamo.exc import UserError, UserErrorType
if dynamic_shapes is None or len(dynamic_shapes) == 0:
# we run with dynamic by default, so no need to produce constraints
return []
if isinstance(dynamic_shapes, (tuple, list)):
combined_args = type(dynamic_shapes)(combined_args.values()) # type: ignore[assignment, misc]
# map of Dim names representing input shape dimensions to constraints on them
symbols: dict[str, list[Constraint]] = defaultdict(list)
# track roots that do not directly represent input shape dimensions
phantom_roots: dict[str, _PhantomRoot] = {}
derived_constraints_with_phantom_root: list[_DerivedConstraint] = []
# list of constraints to return
constraints: list[Constraint] = []
def to_constraint(dim, tensor, i):
import sympy
from torch.fx.experimental.symbolic_shapes import StrictMinMaxConstraint
from torch.utils._sympy.solve import try_solve
from torch.utils._sympy.value_ranges import ValueRanges
def root_value():
# given tensor.shape[i] is the value of dim = fn(root),
# find the value of root
symbol = sympy.Symbol(dim.root.__name__, integer=True)
expr = dim.fn(symbol)
solution = try_solve(sympy.Eq(expr, tensor.shape[i]), symbol)
if solution is not None:
return int(solution[1])
else:
raise UserError( # noqa: B904
UserErrorType.CONSTRAINT_VIOLATION,
f"Expected shape[{i}] = {tensor.shape[i]} of input Tensor to be "
f"of the form {expr}, where {symbol} is an integer",
)
if isinstance(dim, _DerivedDim):
# generate a _DerivedConstraint where the root is:
# - either a _ConstraintTarget (if dim.root directly describes an input shape)
# - or a _PhantomRoot (otherwise)
dim_root = dim.root # type: ignore[attr-defined]
if dim_root.__name__ in symbols:
# root represents an input shape dimension
root_constraint = symbols[dim_root.__name__][0]
root = _ConstraintTarget(
root_constraint.t_id,
root_constraint.dim,
)
elif dim_root.__name__ not in phantom_roots:
# create a phantom root
root = _PhantomRoot( # type: ignore[assignment]
name=dim_root.__name__,
constraint_range=StrictMinMaxConstraint(
vr=ValueRanges(lower=dim_root.min, upper=dim_root.max),
warn_only=False,
),
val=root_value(),
)
phantom_roots[dim_root.__name__] = root # type: ignore[assignment]
else:
root = phantom_roots[dim_root.__name__] # type: ignore[assignment]
constraint = _DerivedConstraint(
id(tensor),
i,
dim.__name__,
StrictMinMaxConstraint(
vr=ValueRanges(lower=dim.min, upper=dim.max),
warn_only=False,
),
root,
dim.fn, # type: ignore[attr-defined]
)
if isinstance(root, _PhantomRoot):
# NOTE(avik): since we have not processed all inputs yet, we may replace this
# with a root that does represent an input shape dimension later (see below)
derived_constraints_with_phantom_root.append(constraint)
elif isinstance(dim, _StaticDim):
constraint = _Constraint( # type: ignore[assignment]
id(tensor),
i,
dim.__name__,
StrictMinMaxConstraint(
vr=ValueRanges(lower=dim.value, upper=dim.value), # type: ignore[attr-defined]
warn_only=False,
),
)
else:
assert isinstance(dim, Dim)
constraint = _Constraint( # type: ignore[assignment]
id(tensor),
i,
dim.__name__,
StrictMinMaxConstraint(
vr=ValueRanges(lower=dim.min, upper=dim.max), # type: ignore[attr-defined]
warn_only=False,
),
)
return constraint
def _parse_tensor_dim(tensor, idx, dim) -> None:
def _create_static_dim(tensor, i, value):
return _StaticDim(value)
if isinstance(dim, (int, Dim)):
if isinstance(dim, int):
dim = _create_static_dim(tensor, idx, dim)
constraint = to_constraint(dim, tensor, idx)
symbols[dim.__name__].append(constraint)
elif isinstance(dim, _DimHint):
if dim.type == _DimHintType.AUTO:
torch._dynamo.maybe_mark_dynamic(tensor, idx)
elif dim.type == _DimHintType.STATIC:
torch._dynamo.mark_static(tensor, idx)
elif dim.type == _DimHintType.DYNAMIC:
torch._dynamo.mark_dynamic(tensor, idx)
constraints.append(_RelaxedConstraint(id(tensor), idx))
elif dim is None:
torch._dynamo.mark_static(tensor, idx)
def update_symbols(path, tensor, shape):
# clean out decorators from user side, or previous export call
# we also delete these attributes in non_strict_utils.py/make_constraints()
tensor._dynamo_weak_dynamic_indices = set()
tensor._dynamo_dynamic_indices = set()
tensor._dynamo_dynamic_range = set()
tensor._dynamo_static_indices = set()
tensor._dynamo_unbacked_indices = set()
if isinstance(shape, dict):
for i, dim in shape.items():
_parse_tensor_dim(tensor, i, dim)
elif isinstance(shape, (tuple, list)):
for i, dim in enumerate(shape):
_parse_tensor_dim(tensor, i, dim)
elif shape is None:
for i in range(tensor.dim()):
_parse_tensor_dim(tensor, i, None)
def assoc_shape(path, t, dynamic_shape):
if isinstance(t, torch.Tensor):
update_symbols(path, t, dynamic_shape)
elif isinstance(t, _IntWrapper):
# If tensor dimensions are marked as dynamic, the tensors themselves
# get marked using mark_dynamic. However since we can't mark
# integers as dynamic, we first wrap integers in this class, and
# then set the `dim` field of the class with the dynamic shapes dim
# to mark the integer as dynamic.
t.dynamism = dynamic_shape
_tree_map_with_path(assoc_shape, combined_args, dynamic_shapes, tree_name="inputs")
for derived_constraint_with_phantom_root in derived_constraints_with_phantom_root:
phantom_root_name = derived_constraint_with_phantom_root.root.name # type: ignore[union-attr]
if phantom_root_name in symbols:
# We found an input shape dimension corresponding to this name, so we
# do not need a phantom symbol for it after all.
# NOTE(avik): Overall we want to maintain the invariant that roots that
# are phantom symbols are really "phantom," i.e., they cannot be represented
# by any input source. This is important when we are deciding derived equalities,
# since we can focus our attention exclusively on input sources: deciding
# derived equalities involving phantom symbols are, in comparison, trivial.
derived_constraint_with_phantom_root.root = symbols[phantom_root_name][0]
for dynamic_dims in symbols.values():
constraints.extend(dynamic_dims)
return constraints
def _get_dim_name_mapping(
dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any], None],
):
name_to_dim = {}
for dim in tree_flatten(
dynamic_shapes,
is_leaf=lambda x: isinstance(x, Dim),
)[0]:
if dim is None:
# NOTE: this must denote a non-Tensor or automatic at this point.
continue
if isinstance(dim, int):
continue
elif isinstance(dim, Dim):
name_to_dim[dim.__name__] = dim
if isinstance(dim, _DerivedDim):
name_to_dim[dim.root.__name__] = dim.root # type: ignore[attr-defined]
else:
assert isinstance(dim, _DimHint)
return name_to_dim
def refine_dynamic_shapes_from_suggested_fixes(
msg: str,
dynamic_shapes: Union[dict[str, Any], tuple[Any], list[Any]],
) -> Union[dict[str, Any], tuple[Any], list[Any]]:
"""
When exporting with :func:`dynamic_shapes`, export may fail with a ConstraintViolation error if the specification
doesn't match the constraints inferred from tracing the model. The error message may provide suggested fixes -
changes that can be made to :func:`dynamic_shapes` to export successfully.
Example ConstraintViolation error message::
Suggested fixes:
dim = Dim('dim', min=3, max=6) # this just refines the dim's range
dim = 4 # this specializes to a constant
dy = dx + 1 # dy was specified as an independent dim, but is actually tied to dx with this relation
This is a helper function that takes the ConstraintViolation error message and the original :func:`dynamic_shapes` spec,
and returns a new :func:`dynamic_shapes` spec that incorporates the suggested fixes.
Example usage::
try:
ep = export(mod, args, dynamic_shapes=dynamic_shapes)
except torch._dynamo.exc.UserError as exc:
new_shapes = refine_dynamic_shapes_from_suggested_fixes(
exc.msg, dynamic_shapes
)
ep = export(mod, args, dynamic_shapes=new_shapes)
"""
import re
import sympy
from torch._dynamo.exc import UserError, UserErrorType
from torch.fx.experimental.symbolic_shapes import _is_supported_equivalence
try:
shape_fixes_msg = msg.split("Suggested fixes:")[1].strip()
except Exception as exc:
raise UserError(
UserErrorType.INVALID_INPUT,
"Suggested fixes not found in error message given to refine_dynamic_shapes_from_suggested_fixes()",
) from exc
# build shape_fixes dictionary
shape_fixes = {}
for fix in shape_fixes_msg.split("\n"):
fix = fix.strip()
if match := re.match(r"(.*) = Dim\('(.*)'.*\)", fix):
name = match.group(1)
_min, _max = None, None
if match_min := re.match(r".* = Dim\('.*', min\=([0-9]+).*\)", fix):
_min = int(match_min.group(1))
if match_max := re.match(r".* = Dim\('.*'.*max\=([0-9]+)\)", fix):
_max = int(match_max.group(1))
shape_fixes[name] = Dim(name, min=_min, max=_max)
else:
name, expr = fix.split(" = ")
expr = sympy.sympify(expr)
if isinstance(expr, sympy.Number):
# static, integer
shape_fixes[name] = int(expr) # type: ignore[assignment]
else:
# relation or derived dim
shape_fixes[name] = expr
name_to_dim = _get_dim_name_mapping(dynamic_shapes)
# track derived dim roots
roots: set[str] = set()
for k, c in shape_fixes.items():
assert isinstance(c, (int, Dim, _DerivedDim, sympy.Expr))
if isinstance(c, sympy.Expr): # check dim/derived dim expression
assert _is_supported_equivalence(c)
shape_fixes[k] = c
roots.add(str(next(iter(c.free_symbols))))
if isinstance(c, _DerivedDim):
roots.add(c.root.__name__) # type: ignore[attr-defined]
# check keys are existing dims or new roots
for k, c in shape_fixes.items():
assert k in name_to_dim or k in roots
# cache so we don't produce multiple derived dim objects
derived_dim_cache: dict[str, _DerivedDim] = {}
def apply_fixes(path, dim, dummy):
if dim is None or isinstance(dim, int): # not dynamic
return dim
elif dim.__name__ in shape_fixes: # directly fix
fix = shape_fixes[dim.__name__]
if isinstance(fix, sympy.Expr): # now derived or related
if str(fix) in derived_dim_cache:
return derived_dim_cache[str(fix)]
else:
symbol = next(iter(fix.free_symbols))
# try to locate symbol
if symbol.name in shape_fixes:
root = shape_fixes[symbol.name]
else:
assert symbol.name in name_to_dim
root = name_to_dim[symbol.name]
# figure out value of fix
modulus, remainder = sympy.polys.polytools.div(fix, symbol)
dim = root
if modulus != 1:
dim = int(modulus) * dim
if remainder != 0:
dim = dim + int(remainder)
derived_dim_cache[str(fix)] = dim
return dim
else:
return fix
elif isinstance(dim, _DerivedDim) and dim.root.__name__ in shape_fixes: # type: ignore[attr-defined]
if dim.__name__ in derived_dim_cache:
return derived_dim_cache[dim.__name__]
else: # evaluate new derived value based on root
_dim = dim.fn(shape_fixes[dim.root.__name__]) # type: ignore[attr-defined]
derived_dim_cache[dim.__name__] = _dim
return _dim
return dim # unchanged dim
return _tree_map_with_path(apply_fixes, dynamic_shapes, dynamic_shapes)
|