_comparison.py 65 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650
  1. # mypy: allow-untyped-defs
  2. import abc
  3. import cmath
  4. import collections.abc
  5. import contextlib
  6. from collections.abc import Callable, Collection, Sequence
  7. from typing import Any, NoReturn, Optional, Union
  8. from typing_extensions import deprecated
  9. import torch
  10. try:
  11. import numpy as np
  12. HAS_NUMPY = True
  13. except ModuleNotFoundError:
  14. HAS_NUMPY = False
  15. np = None # type: ignore[assignment]
  16. class ErrorMeta(Exception):
  17. """Internal testing exception that makes that carries error metadata."""
  18. def __init__(
  19. self, type: type[Exception], msg: str, *, id: tuple[Any, ...] = ()
  20. ) -> None:
  21. super().__init__(
  22. "If you are a user and see this message during normal operation "
  23. "please file an issue at https://github.com/pytorch/pytorch/issues. "
  24. "If you are a developer and working on the comparison functions, please `raise ErrorMeta.to_error()` "
  25. "for user facing errors."
  26. )
  27. self.type = type
  28. self.msg = msg
  29. self.id = id
  30. def to_error(
  31. self, msg: Optional[Union[str, Callable[[str], str]]] = None
  32. ) -> Exception:
  33. if not isinstance(msg, str):
  34. generated_msg = self.msg
  35. if self.id:
  36. generated_msg += f"\n\nThe failure occurred for item {''.join(str([item]) for item in self.id)}"
  37. msg = msg(generated_msg) if callable(msg) else generated_msg
  38. return self.type(msg)
  39. # Some analysis of tolerance by logging tests from test_torch.py can be found in
  40. # https://github.com/pytorch/pytorch/pull/32538.
  41. # {dtype: (rtol, atol)}
  42. _DTYPE_PRECISIONS = {
  43. torch.float16: (0.001, 1e-5),
  44. torch.bfloat16: (0.016, 1e-5),
  45. torch.float32: (1.3e-6, 1e-5),
  46. torch.float64: (1e-7, 1e-7),
  47. torch.complex32: (0.001, 1e-5),
  48. torch.complex64: (1.3e-6, 1e-5),
  49. torch.complex128: (1e-7, 1e-7),
  50. }
  51. # The default tolerances of torch.float32 are used for quantized dtypes, because quantized tensors are compared in
  52. # their dequantized and floating point representation. For more details see `TensorLikePair._compare_quantized_values`
  53. _DTYPE_PRECISIONS.update(
  54. dict.fromkeys(
  55. (torch.quint8, torch.quint2x4, torch.quint4x2, torch.qint8, torch.qint32),
  56. _DTYPE_PRECISIONS[torch.float32],
  57. )
  58. )
  59. def default_tolerances(
  60. *inputs: Union[torch.Tensor, torch.dtype],
  61. dtype_precisions: Optional[dict[torch.dtype, tuple[float, float]]] = None,
  62. ) -> tuple[float, float]:
  63. """Returns the default absolute and relative testing tolerances for a set of inputs based on the dtype.
  64. See :func:`assert_close` for a table of the default tolerance for each dtype.
  65. Returns:
  66. (Tuple[float, float]): Loosest tolerances of all input dtypes.
  67. """
  68. dtypes = []
  69. for input in inputs:
  70. if isinstance(input, torch.Tensor):
  71. dtypes.append(input.dtype)
  72. elif isinstance(input, torch.dtype):
  73. dtypes.append(input)
  74. else:
  75. raise TypeError(
  76. f"Expected a torch.Tensor or a torch.dtype, but got {type(input)} instead."
  77. )
  78. dtype_precisions = dtype_precisions or _DTYPE_PRECISIONS
  79. rtols, atols = zip(
  80. *[dtype_precisions.get(dtype, (0.0, 0.0)) for dtype in dtypes], strict=True
  81. )
  82. return max(rtols), max(atols)
  83. def get_tolerances(
  84. *inputs: Union[torch.Tensor, torch.dtype],
  85. rtol: Optional[float],
  86. atol: Optional[float],
  87. id: tuple[Any, ...] = (),
  88. ) -> tuple[float, float]:
  89. """Gets absolute and relative to be used for numeric comparisons.
  90. If both ``rtol`` and ``atol`` are specified, this is a no-op. If both are not specified, the return value of
  91. :func:`default_tolerances` is used.
  92. Raises:
  93. ErrorMeta: With :class:`ValueError`, if only ``rtol`` or ``atol`` is specified.
  94. Returns:
  95. (Tuple[float, float]): Valid absolute and relative tolerances.
  96. """
  97. if (rtol is None) ^ (atol is None):
  98. # We require both tolerance to be omitted or specified, because specifying only one might lead to surprising
  99. # results. Imagine setting atol=0.0 and the tensors still match because rtol>0.0.
  100. raise ErrorMeta(
  101. ValueError,
  102. f"Both 'rtol' and 'atol' must be either specified or omitted, "
  103. f"but got no {'rtol' if rtol is None else 'atol'}.",
  104. id=id,
  105. )
  106. elif rtol is not None and atol is not None:
  107. return rtol, atol
  108. else:
  109. return default_tolerances(*inputs)
  110. def _make_bitwise_mismatch_msg(
  111. *,
  112. default_identifier: str,
  113. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  114. extra: Optional[str] = None,
  115. first_mismatch_idx: Optional[tuple[int, ...]] = None,
  116. ):
  117. """Makes a mismatch error message for bitwise values.
  118. Args:
  119. default_identifier (str): Default description of the compared values, e.g. "Tensor-likes".
  120. identifier (Optional[Union[str, Callable[[str], str]]]): Optional identifier that overrides
  121. ``default_identifier``. Can be passed as callable in which case it will be called with
  122. ``default_identifier`` to create the description at runtime.
  123. extra (Optional[str]): Extra information to be placed after the message header and the mismatch statistics.
  124. first_mismatch_idx (Optional[tuple[int, ...]]): the index of the first mismatch, for each dimension.
  125. """
  126. if identifier is None:
  127. identifier = default_identifier
  128. elif callable(identifier):
  129. identifier = identifier(default_identifier)
  130. msg = f"{identifier} are not 'equal'!\n\n"
  131. if extra:
  132. msg += f"{extra.strip()}\n"
  133. if first_mismatch_idx is not None:
  134. msg += f"The first mismatched element is at index {first_mismatch_idx}.\n"
  135. return msg.strip()
  136. def _make_mismatch_msg(
  137. *,
  138. default_identifier: str,
  139. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  140. extra: Optional[str] = None,
  141. abs_diff: float,
  142. abs_diff_idx: Optional[Union[int, tuple[int, ...]]] = None,
  143. atol: float,
  144. rel_diff: float,
  145. rel_diff_idx: Optional[Union[int, tuple[int, ...]]] = None,
  146. rtol: float,
  147. ) -> str:
  148. """Makes a mismatch error message for numeric values.
  149. Args:
  150. default_identifier (str): Default description of the compared values, e.g. "Tensor-likes".
  151. identifier (Optional[Union[str, Callable[[str], str]]]): Optional identifier that overrides
  152. ``default_identifier``. Can be passed as callable in which case it will be called with
  153. ``default_identifier`` to create the description at runtime.
  154. extra (Optional[str]): Extra information to be placed after the message header and the mismatch statistics.
  155. abs_diff (float): Absolute difference.
  156. abs_diff_idx (Optional[Union[int, Tuple[int, ...]]]): Optional index of the absolute difference.
  157. atol (float): Allowed absolute tolerance. Will only be added to mismatch statistics if it or ``rtol`` are
  158. ``> 0``.
  159. rel_diff (float): Relative difference.
  160. rel_diff_idx (Optional[Union[int, Tuple[int, ...]]]): Optional index of the relative difference.
  161. rtol (float): Allowed relative tolerance. Will only be added to mismatch statistics if it or ``atol`` are
  162. ``> 0``.
  163. """
  164. equality = rtol == 0 and atol == 0
  165. def make_diff_msg(
  166. *,
  167. type: str,
  168. diff: float,
  169. idx: Optional[Union[int, tuple[int, ...]]],
  170. tol: float,
  171. ) -> str:
  172. if idx is None:
  173. msg = f"{type.title()} difference: {diff}"
  174. else:
  175. msg = f"Greatest {type} difference: {diff} at index {idx}"
  176. if not equality:
  177. msg += f" (up to {tol} allowed)"
  178. return msg + "\n"
  179. if identifier is None:
  180. identifier = default_identifier
  181. elif callable(identifier):
  182. identifier = identifier(default_identifier)
  183. msg = f"{identifier} are not {'equal' if equality else 'close'}!\n\n"
  184. if extra:
  185. msg += f"{extra.strip()}\n"
  186. msg += make_diff_msg(type="absolute", diff=abs_diff, idx=abs_diff_idx, tol=atol)
  187. msg += make_diff_msg(type="relative", diff=rel_diff, idx=rel_diff_idx, tol=rtol)
  188. return msg.strip()
  189. def make_scalar_mismatch_msg(
  190. actual: Union[bool, int, float, complex],
  191. expected: Union[bool, int, float, complex],
  192. *,
  193. rtol: float,
  194. atol: float,
  195. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  196. ) -> str:
  197. """Makes a mismatch error message for scalars.
  198. Args:
  199. actual (Union[bool, int, float, complex]): Actual scalar.
  200. expected (Union[bool, int, float, complex]): Expected scalar.
  201. rtol (float): Relative tolerance.
  202. atol (float): Absolute tolerance.
  203. identifier (Optional[Union[str, Callable[[str], str]]]): Optional description for the scalars. Can be passed
  204. as callable in which case it will be called by the default value to create the description at runtime.
  205. Defaults to "Scalars".
  206. """
  207. abs_diff = abs(actual - expected)
  208. # pyrefly: ignore [bad-argument-type]
  209. rel_diff = float("inf") if expected == 0 else abs_diff / abs(expected)
  210. return _make_mismatch_msg(
  211. default_identifier="Scalars",
  212. identifier=identifier,
  213. extra=f"Expected {expected} but got {actual}.",
  214. abs_diff=abs_diff,
  215. atol=atol,
  216. rel_diff=rel_diff,
  217. rtol=rtol,
  218. )
  219. def make_tensor_mismatch_msg(
  220. actual: torch.Tensor,
  221. expected: torch.Tensor,
  222. matches: torch.Tensor,
  223. *,
  224. rtol: float,
  225. atol: float,
  226. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  227. ):
  228. """Makes a mismatch error message for tensors.
  229. Args:
  230. actual (torch.Tensor): Actual tensor.
  231. expected (torch.Tensor): Expected tensor.
  232. matches (torch.Tensor): Boolean mask of the same shape as ``actual`` and ``expected`` that indicates the
  233. location of matches.
  234. rtol (float): Relative tolerance.
  235. atol (float): Absolute tolerance.
  236. identifier (Optional[Union[str, Callable[[str], str]]]): Optional description for the tensors. Can be passed
  237. as callable in which case it will be called by the default value to create the description at runtime.
  238. Defaults to "Tensor-likes".
  239. """
  240. def unravel_flat_index(flat_index: int) -> tuple[int, ...]:
  241. if not matches.shape:
  242. return ()
  243. inverse_index = []
  244. for size in matches.shape[::-1]:
  245. div, mod = divmod(flat_index, size)
  246. flat_index = div
  247. inverse_index.append(mod)
  248. return tuple(inverse_index[::-1])
  249. number_of_elements = matches.numel()
  250. total_mismatches = number_of_elements - int(torch.sum(matches))
  251. extra = (
  252. f"Mismatched elements: {total_mismatches} / {number_of_elements} "
  253. f"({total_mismatches / number_of_elements:.1%})"
  254. )
  255. if actual.dtype.is_floating_point and actual.dtype.itemsize == 1:
  256. # skip checking for max_abs_diff and max_rel_diff for float8-like values
  257. first_mismatch_idx = tuple(torch.nonzero(~matches, as_tuple=False)[0].tolist())
  258. return _make_bitwise_mismatch_msg(
  259. default_identifier="Tensor-likes",
  260. identifier=identifier,
  261. extra=extra,
  262. first_mismatch_idx=first_mismatch_idx,
  263. )
  264. actual_flat = actual.flatten()
  265. expected_flat = expected.flatten()
  266. matches_flat = matches.flatten()
  267. if not actual.dtype.is_floating_point and not actual.dtype.is_complex:
  268. # TODO: Instead of always upcasting to int64, it would be sufficient to cast to the next higher dtype to avoid
  269. # overflow
  270. actual_flat = actual_flat.to(torch.int64)
  271. expected_flat = expected_flat.to(torch.int64)
  272. abs_diff = torch.abs(actual_flat - expected_flat)
  273. # Ensure that only mismatches are used for the max_abs_diff computation
  274. abs_diff[matches_flat] = 0
  275. max_abs_diff, max_abs_diff_flat_idx = torch.max(abs_diff, 0)
  276. rel_diff = abs_diff / torch.abs(expected_flat)
  277. # Ensure that only mismatches are used for the max_rel_diff computation
  278. rel_diff[matches_flat] = 0
  279. max_rel_diff, max_rel_diff_flat_idx = torch.max(rel_diff, 0)
  280. return _make_mismatch_msg(
  281. default_identifier="Tensor-likes",
  282. identifier=identifier,
  283. extra=extra,
  284. abs_diff=max_abs_diff.item(),
  285. abs_diff_idx=unravel_flat_index(int(max_abs_diff_flat_idx)),
  286. atol=atol,
  287. rel_diff=max_rel_diff.item(),
  288. rel_diff_idx=unravel_flat_index(int(max_rel_diff_flat_idx)),
  289. rtol=rtol,
  290. )
  291. class UnsupportedInputs(Exception): # noqa: B903
  292. """Exception to be raised during the construction of a :class:`Pair` in case it doesn't support the inputs."""
  293. class Pair(abc.ABC):
  294. """ABC for all comparison pairs to be used in conjunction with :func:`assert_equal`.
  295. Each subclass needs to overwrite :meth:`Pair.compare` that performs the actual comparison.
  296. Each pair receives **all** options, so select the ones applicable for the subclass and forward the rest to the
  297. super class. Raising an :class:`UnsupportedInputs` during constructions indicates that the pair is not able to
  298. handle the inputs and the next pair type will be tried.
  299. All other errors should be raised as :class:`ErrorMeta`. After the instantiation, :meth:`Pair._make_error_meta` can
  300. be used to automatically handle overwriting the message with a user supplied one and id handling.
  301. """
  302. def __init__(
  303. self,
  304. actual: Any,
  305. expected: Any,
  306. *,
  307. id: tuple[Any, ...] = (),
  308. **unknown_parameters: Any,
  309. ) -> None:
  310. self.actual = actual
  311. self.expected = expected
  312. self.id = id
  313. self._unknown_parameters = unknown_parameters
  314. @staticmethod
  315. def _inputs_not_supported() -> NoReturn:
  316. raise UnsupportedInputs
  317. @staticmethod
  318. def _check_inputs_isinstance(*inputs: Any, cls: Union[type, tuple[type, ...]]):
  319. """Checks if all inputs are instances of a given class and raise :class:`UnsupportedInputs` otherwise."""
  320. if not all(isinstance(input, cls) for input in inputs):
  321. Pair._inputs_not_supported()
  322. def _fail(
  323. self, type: type[Exception], msg: str, *, id: tuple[Any, ...] = ()
  324. ) -> NoReturn:
  325. """Raises an :class:`ErrorMeta` from a given exception type and message and the stored id.
  326. .. warning::
  327. If you use this before the ``super().__init__(...)`` call in the constructor, you have to pass the ``id``
  328. explicitly.
  329. """
  330. raise ErrorMeta(type, msg, id=self.id if not id and hasattr(self, "id") else id)
  331. @abc.abstractmethod
  332. def compare(self) -> None:
  333. """Compares the inputs and raises an :class`ErrorMeta` in case they mismatch."""
  334. def extra_repr(self) -> Sequence[Union[str, tuple[str, Any]]]:
  335. """Returns extra information that will be included in the representation.
  336. Should be overwritten by all subclasses that use additional options. The representation of the object will only
  337. be surfaced in case we encounter an unexpected error and thus should help debug the issue. Can be a sequence of
  338. key-value-pairs or attribute names.
  339. """
  340. return []
  341. def __repr__(self) -> str:
  342. head = f"{type(self).__name__}("
  343. tail = ")"
  344. body = [
  345. f" {name}={value!s},"
  346. for name, value in [
  347. ("id", self.id),
  348. ("actual", self.actual),
  349. ("expected", self.expected),
  350. *[
  351. (extra, getattr(self, extra)) if isinstance(extra, str) else extra
  352. for extra in self.extra_repr()
  353. ],
  354. ]
  355. ]
  356. return "\n".join((head, *body, *tail))
  357. class ObjectPair(Pair):
  358. """Pair for any type of inputs that will be compared with the `==` operator.
  359. .. note::
  360. Since this will instantiate for any kind of inputs, it should only be used as fallback after all other pairs
  361. couldn't handle the inputs.
  362. """
  363. def compare(self) -> None:
  364. try:
  365. equal = self.actual == self.expected
  366. except Exception as error:
  367. # We are not using `self._raise_error_meta` here since we need the exception chaining
  368. raise ErrorMeta(
  369. ValueError,
  370. f"{self.actual} == {self.expected} failed with:\n{error}.",
  371. id=self.id,
  372. ) from error
  373. if not equal:
  374. self._fail(AssertionError, f"{self.actual} != {self.expected}")
  375. class NonePair(Pair):
  376. """Pair for ``None`` inputs."""
  377. def __init__(self, actual: Any, expected: Any, **other_parameters: Any) -> None:
  378. if not (actual is None or expected is None):
  379. self._inputs_not_supported()
  380. super().__init__(actual, expected, **other_parameters)
  381. def compare(self) -> None:
  382. if not (self.actual is None and self.expected is None):
  383. self._fail(
  384. AssertionError, f"None mismatch: {self.actual} is not {self.expected}"
  385. )
  386. class BooleanPair(Pair):
  387. """Pair for :class:`bool` inputs.
  388. .. note::
  389. If ``numpy`` is available, also handles :class:`numpy.bool_` inputs.
  390. """
  391. def __init__(
  392. self,
  393. actual: Any,
  394. expected: Any,
  395. *,
  396. id: tuple[Any, ...],
  397. **other_parameters: Any,
  398. ) -> None:
  399. actual, expected = self._process_inputs(actual, expected, id=id)
  400. super().__init__(actual, expected, **other_parameters)
  401. @property
  402. def _supported_types(self) -> tuple[type, ...]:
  403. cls: list[type] = [bool]
  404. if HAS_NUMPY:
  405. # pyrefly: ignore [missing-attribute]
  406. cls.append(np.bool_)
  407. return tuple(cls)
  408. def _process_inputs(
  409. self, actual: Any, expected: Any, *, id: tuple[Any, ...]
  410. ) -> tuple[bool, bool]:
  411. self._check_inputs_isinstance(actual, expected, cls=self._supported_types)
  412. actual, expected = (
  413. self._to_bool(bool_like, id=id) for bool_like in (actual, expected)
  414. )
  415. return actual, expected
  416. def _to_bool(self, bool_like: Any, *, id: tuple[Any, ...]) -> bool:
  417. if isinstance(bool_like, bool):
  418. return bool_like
  419. # pyrefly: ignore [missing-attribute]
  420. elif isinstance(bool_like, np.bool_):
  421. return bool_like.item()
  422. else:
  423. raise ErrorMeta(
  424. TypeError, f"Unknown boolean type {type(bool_like)}.", id=id
  425. )
  426. def compare(self) -> None:
  427. if self.actual is not self.expected:
  428. self._fail(
  429. AssertionError,
  430. f"Booleans mismatch: {self.actual} is not {self.expected}",
  431. )
  432. class NumberPair(Pair):
  433. """Pair for Python number (:class:`int`, :class:`float`, and :class:`complex`) inputs.
  434. .. note::
  435. If ``numpy`` is available, also handles :class:`numpy.number` inputs.
  436. Kwargs:
  437. rtol (Optional[float]): Relative tolerance. If specified ``atol`` must also be specified. If omitted, default
  438. values based on the type are selected with the below table.
  439. atol (Optional[float]): Absolute tolerance. If specified ``rtol`` must also be specified. If omitted, default
  440. values based on the type are selected with the below table.
  441. equal_nan (bool): If ``True``, two ``NaN`` values are considered equal. Defaults to ``False``.
  442. check_dtype (bool): If ``True``, the type of the inputs will be checked for equality. Defaults to ``False``.
  443. The following table displays correspondence between Python number type and the ``torch.dtype``'s. See
  444. :func:`assert_close` for the corresponding tolerances.
  445. +------------------+-------------------------------+
  446. | ``type`` | corresponding ``torch.dtype`` |
  447. +==================+===============================+
  448. | :class:`int` | :attr:`~torch.int64` |
  449. +------------------+-------------------------------+
  450. | :class:`float` | :attr:`~torch.float64` |
  451. +------------------+-------------------------------+
  452. | :class:`complex` | :attr:`~torch.complex64` |
  453. +------------------+-------------------------------+
  454. """
  455. _TYPE_TO_DTYPE = {
  456. int: torch.int64,
  457. float: torch.float64,
  458. complex: torch.complex128,
  459. }
  460. _NUMBER_TYPES = tuple(_TYPE_TO_DTYPE.keys())
  461. def __init__(
  462. self,
  463. actual: Any,
  464. expected: Any,
  465. *,
  466. id: tuple[Any, ...] = (),
  467. rtol: Optional[float] = None,
  468. atol: Optional[float] = None,
  469. equal_nan: bool = False,
  470. check_dtype: bool = False,
  471. **other_parameters: Any,
  472. ) -> None:
  473. actual, expected = self._process_inputs(actual, expected, id=id)
  474. super().__init__(actual, expected, id=id, **other_parameters)
  475. self.rtol, self.atol = get_tolerances(
  476. *[self._TYPE_TO_DTYPE[type(input)] for input in (actual, expected)],
  477. rtol=rtol,
  478. atol=atol,
  479. id=id,
  480. )
  481. self.equal_nan = equal_nan
  482. self.check_dtype = check_dtype
  483. @property
  484. def _supported_types(self) -> tuple[type, ...]:
  485. cls = list(self._NUMBER_TYPES)
  486. if HAS_NUMPY:
  487. # pyrefly: ignore [missing-attribute]
  488. cls.append(np.number)
  489. return tuple(cls)
  490. def _process_inputs(
  491. self, actual: Any, expected: Any, *, id: tuple[Any, ...]
  492. ) -> tuple[Union[int, float, complex], Union[int, float, complex]]:
  493. self._check_inputs_isinstance(actual, expected, cls=self._supported_types)
  494. actual, expected = (
  495. self._to_number(number_like, id=id) for number_like in (actual, expected)
  496. )
  497. return actual, expected
  498. def _to_number(
  499. self, number_like: Any, *, id: tuple[Any, ...]
  500. ) -> Union[int, float, complex]:
  501. # pyrefly: ignore [missing-attribute]
  502. if HAS_NUMPY and isinstance(number_like, np.number):
  503. return number_like.item()
  504. elif isinstance(number_like, self._NUMBER_TYPES):
  505. return number_like # type: ignore[return-value]
  506. else:
  507. raise ErrorMeta(
  508. TypeError, f"Unknown number type {type(number_like)}.", id=id
  509. )
  510. def compare(self) -> None:
  511. if self.check_dtype and type(self.actual) is not type(self.expected):
  512. self._fail(
  513. AssertionError,
  514. f"The (d)types do not match: {type(self.actual)} != {type(self.expected)}.",
  515. )
  516. if self.actual == self.expected:
  517. return
  518. if self.equal_nan and cmath.isnan(self.actual) and cmath.isnan(self.expected):
  519. return
  520. abs_diff = abs(self.actual - self.expected)
  521. tolerance = self.atol + self.rtol * abs(self.expected)
  522. if cmath.isfinite(abs_diff) and abs_diff <= tolerance:
  523. return
  524. self._fail(
  525. AssertionError,
  526. make_scalar_mismatch_msg(
  527. self.actual, self.expected, rtol=self.rtol, atol=self.atol
  528. ),
  529. )
  530. def extra_repr(self) -> Sequence[str]:
  531. return (
  532. "rtol",
  533. "atol",
  534. "equal_nan",
  535. "check_dtype",
  536. )
  537. class TensorLikePair(Pair):
  538. """Pair for :class:`torch.Tensor`-like inputs.
  539. Kwargs:
  540. allow_subclasses (bool):
  541. rtol (Optional[float]): Relative tolerance. If specified ``atol`` must also be specified. If omitted, default
  542. values based on the type are selected. See :func:assert_close: for details.
  543. atol (Optional[float]): Absolute tolerance. If specified ``rtol`` must also be specified. If omitted, default
  544. values based on the type are selected. See :func:assert_close: for details.
  545. equal_nan (bool): If ``True``, two ``NaN`` values are considered equal. Defaults to ``False``.
  546. check_device (bool): If ``True`` (default), asserts that corresponding tensors are on the same
  547. :attr:`~torch.Tensor.device`. If this check is disabled, tensors on different
  548. :attr:`~torch.Tensor.device`'s are moved to the CPU before being compared.
  549. check_dtype (bool): If ``True`` (default), asserts that corresponding tensors have the same ``dtype``. If this
  550. check is disabled, tensors with different ``dtype``'s are promoted to a common ``dtype`` (according to
  551. :func:`torch.promote_types`) before being compared.
  552. check_layout (bool): If ``True`` (default), asserts that corresponding tensors have the same ``layout``. If this
  553. check is disabled, tensors with different ``layout``'s are converted to strided tensors before being
  554. compared.
  555. check_stride (bool): If ``True`` and corresponding tensors are strided, asserts that they have the same stride.
  556. """
  557. def __init__(
  558. self,
  559. actual: Any,
  560. expected: Any,
  561. *,
  562. id: tuple[Any, ...] = (),
  563. allow_subclasses: bool = True,
  564. rtol: Optional[float] = None,
  565. atol: Optional[float] = None,
  566. equal_nan: bool = False,
  567. check_device: bool = True,
  568. check_dtype: bool = True,
  569. check_layout: bool = True,
  570. check_stride: bool = False,
  571. **other_parameters: Any,
  572. ):
  573. actual, expected = self._process_inputs(
  574. actual, expected, id=id, allow_subclasses=allow_subclasses
  575. )
  576. super().__init__(actual, expected, id=id, **other_parameters)
  577. self.rtol, self.atol = get_tolerances(
  578. actual, expected, rtol=rtol, atol=atol, id=self.id
  579. )
  580. self.equal_nan = equal_nan
  581. self.check_device = check_device
  582. self.check_dtype = check_dtype
  583. self.check_layout = check_layout
  584. self.check_stride = check_stride
  585. def _process_inputs(
  586. self, actual: Any, expected: Any, *, id: tuple[Any, ...], allow_subclasses: bool
  587. ) -> tuple[torch.Tensor, torch.Tensor]:
  588. directly_related = isinstance(actual, type(expected)) or isinstance(
  589. expected, type(actual)
  590. )
  591. if not directly_related:
  592. self._inputs_not_supported()
  593. if not allow_subclasses and type(actual) is not type(expected):
  594. self._inputs_not_supported()
  595. actual, expected = (self._to_tensor(input) for input in (actual, expected))
  596. for tensor in (actual, expected):
  597. self._check_supported(tensor, id=id)
  598. return actual, expected
  599. def _to_tensor(self, tensor_like: Any) -> torch.Tensor:
  600. if isinstance(tensor_like, torch.Tensor):
  601. return tensor_like
  602. try:
  603. return torch.as_tensor(tensor_like)
  604. except Exception:
  605. self._inputs_not_supported()
  606. def _check_supported(self, tensor: torch.Tensor, *, id: tuple[Any, ...]) -> None:
  607. if tensor.layout not in {
  608. torch.strided,
  609. torch.jagged,
  610. torch.sparse_coo,
  611. torch.sparse_csr,
  612. torch.sparse_csc,
  613. torch.sparse_bsr,
  614. torch.sparse_bsc,
  615. }:
  616. raise ErrorMeta(
  617. ValueError, f"Unsupported tensor layout {tensor.layout}", id=id
  618. )
  619. def compare(self) -> None:
  620. actual, expected = self.actual, self.expected
  621. self._compare_attributes(actual, expected)
  622. if any(input.device.type == "meta" for input in (actual, expected)):
  623. return
  624. actual, expected = self._equalize_attributes(actual, expected)
  625. self._compare_values(actual, expected)
  626. def _compare_attributes(
  627. self,
  628. actual: torch.Tensor,
  629. expected: torch.Tensor,
  630. ) -> None:
  631. """Checks if the attributes of two tensors match.
  632. Always checks
  633. - the :attr:`~torch.Tensor.shape`,
  634. - whether both inputs are quantized or not,
  635. - and if they use the same quantization scheme.
  636. Checks for
  637. - :attr:`~torch.Tensor.layout`,
  638. - :meth:`~torch.Tensor.stride`,
  639. - :attr:`~torch.Tensor.device`, and
  640. - :attr:`~torch.Tensor.dtype`
  641. are optional and can be disabled through the corresponding ``check_*`` flag during construction of the pair.
  642. """
  643. def raise_mismatch_error(
  644. attribute_name: str, actual_value: Any, expected_value: Any
  645. ) -> NoReturn:
  646. self._fail(
  647. AssertionError,
  648. f"The values for attribute '{attribute_name}' do not match: {actual_value} != {expected_value}.",
  649. )
  650. if actual.shape != expected.shape:
  651. raise_mismatch_error("shape", actual.shape, expected.shape)
  652. if actual.is_quantized != expected.is_quantized:
  653. raise_mismatch_error(
  654. "is_quantized", actual.is_quantized, expected.is_quantized
  655. )
  656. elif actual.is_quantized and actual.qscheme() != expected.qscheme():
  657. raise_mismatch_error("qscheme()", actual.qscheme(), expected.qscheme())
  658. if actual.layout != expected.layout:
  659. if self.check_layout:
  660. raise_mismatch_error("layout", actual.layout, expected.layout)
  661. elif (
  662. actual.layout == torch.strided
  663. and self.check_stride
  664. and actual.stride() != expected.stride()
  665. ):
  666. raise_mismatch_error("stride()", actual.stride(), expected.stride())
  667. if self.check_device and actual.device != expected.device:
  668. raise_mismatch_error("device", actual.device, expected.device)
  669. if self.check_dtype and actual.dtype != expected.dtype:
  670. raise_mismatch_error("dtype", actual.dtype, expected.dtype)
  671. def _equalize_attributes(
  672. self, actual: torch.Tensor, expected: torch.Tensor
  673. ) -> tuple[torch.Tensor, torch.Tensor]:
  674. """Equalizes some attributes of two tensors for value comparison.
  675. If ``actual`` and ``expected`` are ...
  676. - ... not on the same :attr:`~torch.Tensor.device`, they are moved CPU memory.
  677. - ... not of the same ``dtype``, they are promoted to a common ``dtype`` (according to
  678. :func:`torch.promote_types`).
  679. - ... not of the same ``layout``, they are converted to strided tensors.
  680. Args:
  681. actual (Tensor): Actual tensor.
  682. expected (Tensor): Expected tensor.
  683. Returns:
  684. (Tuple[Tensor, Tensor]): Equalized tensors.
  685. """
  686. # The comparison logic uses operators currently not supported by the MPS backends.
  687. # See https://github.com/pytorch/pytorch/issues/77144 for details.
  688. # TODO: Remove this conversion as soon as all operations are supported natively by the MPS backend
  689. if actual.is_mps or expected.is_mps: # type: ignore[attr-defined]
  690. actual = actual.cpu()
  691. expected = expected.cpu()
  692. if actual.device != expected.device:
  693. actual = actual.cpu()
  694. expected = expected.cpu()
  695. if actual.dtype != expected.dtype:
  696. actual_dtype = actual.dtype
  697. expected_dtype = expected.dtype
  698. # For uint64, this is not sound in general, which is why promote_types doesn't
  699. # allow it, but for easy testing, we're unlikely to get confused
  700. # by large uint64 overflowing into negative int64
  701. if actual_dtype in [torch.uint64, torch.uint32, torch.uint16]:
  702. actual_dtype = torch.int64
  703. if expected_dtype in [torch.uint64, torch.uint32, torch.uint16]:
  704. expected_dtype = torch.int64
  705. dtype = torch.promote_types(actual_dtype, expected_dtype)
  706. actual = actual.to(dtype)
  707. expected = expected.to(dtype)
  708. if actual.layout != expected.layout:
  709. # These checks are needed, since Tensor.to_dense() fails on tensors that are already strided
  710. actual = actual.to_dense() if actual.layout != torch.strided else actual
  711. expected = (
  712. expected.to_dense() if expected.layout != torch.strided else expected
  713. )
  714. return actual, expected
  715. def _compare_values(self, actual: torch.Tensor, expected: torch.Tensor) -> None:
  716. if actual.is_quantized:
  717. compare_fn = self._compare_quantized_values
  718. elif actual.is_sparse:
  719. compare_fn = self._compare_sparse_coo_values
  720. elif actual.layout in {
  721. torch.sparse_csr,
  722. torch.sparse_csc,
  723. torch.sparse_bsr,
  724. torch.sparse_bsc,
  725. }:
  726. compare_fn = self._compare_sparse_compressed_values
  727. elif actual.layout == torch.jagged:
  728. actual, expected = actual.values(), expected.values()
  729. compare_fn = self._compare_regular_values_close
  730. elif actual.dtype.is_floating_point and actual.dtype.itemsize == 1:
  731. def bitwise_comp(
  732. actual: torch.Tensor,
  733. expected: torch.Tensor,
  734. *,
  735. rtol: float,
  736. atol: float,
  737. equal_nan: bool,
  738. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  739. ) -> None:
  740. if rtol != 0.0 or atol != 0.0:
  741. raise ErrorMeta(
  742. AssertionError,
  743. f"Rtol={rtol} and atol={atol} are not supported for bitwise comparison of low"
  744. " dimensional floats. Please use rtol=0.0 and atol=0.0.",
  745. )
  746. return self._compare_regular_values_close(
  747. actual,
  748. expected,
  749. rtol=rtol,
  750. atol=atol,
  751. equal_nan=equal_nan,
  752. identifier=identifier,
  753. )
  754. compare_fn = bitwise_comp
  755. else:
  756. compare_fn = self._compare_regular_values_close
  757. compare_fn(
  758. actual, expected, rtol=self.rtol, atol=self.atol, equal_nan=self.equal_nan
  759. )
  760. def _compare_quantized_values(
  761. self,
  762. actual: torch.Tensor,
  763. expected: torch.Tensor,
  764. *,
  765. rtol: float,
  766. atol: float,
  767. equal_nan: bool,
  768. ) -> None:
  769. """Compares quantized tensors by comparing the :meth:`~torch.Tensor.dequantize`'d variants for closeness.
  770. .. note::
  771. A detailed discussion about why only the dequantized variant is checked for closeness rather than checking
  772. the individual quantization parameters for closeness and the integer representation for equality can be
  773. found in https://github.com/pytorch/pytorch/issues/68548.
  774. """
  775. return self._compare_regular_values_close(
  776. actual.dequantize(),
  777. expected.dequantize(),
  778. rtol=rtol,
  779. atol=atol,
  780. equal_nan=equal_nan,
  781. identifier=lambda default_identifier: f"Quantized {default_identifier.lower()}",
  782. )
  783. def _compare_sparse_coo_values(
  784. self,
  785. actual: torch.Tensor,
  786. expected: torch.Tensor,
  787. *,
  788. rtol: float,
  789. atol: float,
  790. equal_nan: bool,
  791. ) -> None:
  792. """Compares sparse COO tensors by comparing
  793. - the number of sparse dimensions,
  794. - the number of non-zero elements (nnz) for equality,
  795. - the indices for equality, and
  796. - the values for closeness.
  797. """
  798. if actual.sparse_dim() != expected.sparse_dim():
  799. self._fail(
  800. AssertionError,
  801. (
  802. f"The number of sparse dimensions in sparse COO tensors does not match: "
  803. f"{actual.sparse_dim()} != {expected.sparse_dim()}"
  804. ),
  805. )
  806. if actual._nnz() != expected._nnz():
  807. self._fail(
  808. AssertionError,
  809. (
  810. f"The number of specified values in sparse COO tensors does not match: "
  811. f"{actual._nnz()} != {expected._nnz()}"
  812. ),
  813. )
  814. self._compare_regular_values_equal(
  815. actual._indices(),
  816. expected._indices(),
  817. identifier="Sparse COO indices",
  818. )
  819. self._compare_regular_values_close(
  820. actual._values(),
  821. expected._values(),
  822. rtol=rtol,
  823. atol=atol,
  824. equal_nan=equal_nan,
  825. identifier="Sparse COO values",
  826. )
  827. def _compare_sparse_compressed_values(
  828. self,
  829. actual: torch.Tensor,
  830. expected: torch.Tensor,
  831. *,
  832. rtol: float,
  833. atol: float,
  834. equal_nan: bool,
  835. ) -> None:
  836. """Compares sparse compressed tensors by comparing
  837. - the number of non-zero elements (nnz) for equality,
  838. - the plain indices for equality,
  839. - the compressed indices for equality, and
  840. - the values for closeness.
  841. """
  842. format_name, compressed_indices_method, plain_indices_method = {
  843. torch.sparse_csr: (
  844. "CSR",
  845. torch.Tensor.crow_indices,
  846. torch.Tensor.col_indices,
  847. ),
  848. torch.sparse_csc: (
  849. "CSC",
  850. torch.Tensor.ccol_indices,
  851. torch.Tensor.row_indices,
  852. ),
  853. torch.sparse_bsr: (
  854. "BSR",
  855. torch.Tensor.crow_indices,
  856. torch.Tensor.col_indices,
  857. ),
  858. torch.sparse_bsc: (
  859. "BSC",
  860. torch.Tensor.ccol_indices,
  861. torch.Tensor.row_indices,
  862. ),
  863. }[actual.layout]
  864. if actual._nnz() != expected._nnz():
  865. self._fail(
  866. AssertionError,
  867. (
  868. f"The number of specified values in sparse {format_name} tensors does not match: "
  869. f"{actual._nnz()} != {expected._nnz()}"
  870. ),
  871. )
  872. # Compressed and plain indices in the CSR / CSC / BSR / BSC sparse formats can be `torch.int32` _or_
  873. # `torch.int64`. While the same dtype is enforced for the compressed and plain indices of a single tensor, it
  874. # can be different between two tensors. Thus, we need to convert them to the same dtype, or the comparison will
  875. # fail.
  876. actual_compressed_indices = compressed_indices_method(actual)
  877. expected_compressed_indices = compressed_indices_method(expected)
  878. indices_dtype = torch.promote_types(
  879. actual_compressed_indices.dtype, expected_compressed_indices.dtype
  880. )
  881. self._compare_regular_values_equal(
  882. actual_compressed_indices.to(indices_dtype),
  883. expected_compressed_indices.to(indices_dtype),
  884. identifier=f"Sparse {format_name} {compressed_indices_method.__name__}",
  885. )
  886. self._compare_regular_values_equal(
  887. plain_indices_method(actual).to(indices_dtype),
  888. plain_indices_method(expected).to(indices_dtype),
  889. identifier=f"Sparse {format_name} {plain_indices_method.__name__}",
  890. )
  891. self._compare_regular_values_close(
  892. actual.values(),
  893. expected.values(),
  894. rtol=rtol,
  895. atol=atol,
  896. equal_nan=equal_nan,
  897. identifier=f"Sparse {format_name} values",
  898. )
  899. def _compare_regular_values_equal(
  900. self,
  901. actual: torch.Tensor,
  902. expected: torch.Tensor,
  903. *,
  904. equal_nan: bool = False,
  905. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  906. ) -> None:
  907. """Checks if the values of two tensors are equal."""
  908. self._compare_regular_values_close(
  909. actual, expected, rtol=0, atol=0, equal_nan=equal_nan, identifier=identifier
  910. )
  911. def _compare_regular_values_close(
  912. self,
  913. actual: torch.Tensor,
  914. expected: torch.Tensor,
  915. *,
  916. rtol: float,
  917. atol: float,
  918. equal_nan: bool,
  919. identifier: Optional[Union[str, Callable[[str], str]]] = None,
  920. ) -> None:
  921. """Checks if the values of two tensors are close up to a desired tolerance."""
  922. matches = torch.isclose(
  923. actual, expected, rtol=rtol, atol=atol, equal_nan=equal_nan
  924. )
  925. if torch.all(matches):
  926. return
  927. if actual.shape == torch.Size([]):
  928. msg = make_scalar_mismatch_msg(
  929. actual.item(),
  930. expected.item(),
  931. rtol=rtol,
  932. atol=atol,
  933. identifier=identifier,
  934. )
  935. else:
  936. msg = make_tensor_mismatch_msg(
  937. actual, expected, matches, rtol=rtol, atol=atol, identifier=identifier
  938. )
  939. self._fail(AssertionError, msg)
  940. def extra_repr(self) -> Sequence[str]:
  941. return (
  942. "rtol",
  943. "atol",
  944. "equal_nan",
  945. "check_device",
  946. "check_dtype",
  947. "check_layout",
  948. "check_stride",
  949. )
  950. def originate_pairs(
  951. actual: Any,
  952. expected: Any,
  953. *,
  954. pair_types: Sequence[type[Pair]],
  955. sequence_types: tuple[type, ...] = (collections.abc.Sequence,),
  956. mapping_types: tuple[type, ...] = (collections.abc.Mapping,),
  957. id: tuple[Any, ...] = (),
  958. **options: Any,
  959. # pyrefly: ignore [bad-return]
  960. ) -> list[Pair]:
  961. """Originates pairs from the individual inputs.
  962. ``actual`` and ``expected`` can be possibly nested :class:`~collections.abc.Sequence`'s or
  963. :class:`~collections.abc.Mapping`'s. In this case the pairs are originated by recursing through them.
  964. Args:
  965. actual (Any): Actual input.
  966. expected (Any): Expected input.
  967. pair_types (Sequence[Type[Pair]]): Sequence of pair types that will be tried to construct with the inputs.
  968. First successful pair will be used.
  969. sequence_types (Tuple[Type, ...]): Optional types treated as sequences that will be checked elementwise.
  970. mapping_types (Tuple[Type, ...]): Optional types treated as mappings that will be checked elementwise.
  971. id (Tuple[Any, ...]): Optional id of a pair that will be included in an error message.
  972. **options (Any): Options passed to each pair during construction.
  973. Raises:
  974. ErrorMeta: With :class`AssertionError`, if the inputs are :class:`~collections.abc.Sequence`'s, but their
  975. length does not match.
  976. ErrorMeta: With :class`AssertionError`, if the inputs are :class:`~collections.abc.Mapping`'s, but their set of
  977. keys do not match.
  978. ErrorMeta: With :class`TypeError`, if no pair is able to handle the inputs.
  979. ErrorMeta: With any expected exception that happens during the construction of a pair.
  980. Returns:
  981. (List[Pair]): Originated pairs.
  982. """
  983. # We explicitly exclude str's here since they are self-referential and would cause an infinite recursion loop:
  984. # "a" == "a"[0][0]...
  985. if (
  986. isinstance(actual, sequence_types)
  987. and not isinstance(actual, str)
  988. and isinstance(expected, sequence_types)
  989. and not isinstance(expected, str)
  990. ):
  991. actual_len = len(actual) # type: ignore[arg-type]
  992. expected_len = len(expected) # type: ignore[arg-type]
  993. if actual_len != expected_len:
  994. raise ErrorMeta(
  995. AssertionError,
  996. f"The length of the sequences mismatch: {actual_len} != {expected_len}",
  997. id=id,
  998. )
  999. pairs = []
  1000. for idx in range(actual_len):
  1001. pairs.extend(
  1002. originate_pairs(
  1003. actual[idx], # type: ignore[index]
  1004. expected[idx], # type: ignore[index]
  1005. pair_types=pair_types,
  1006. sequence_types=sequence_types,
  1007. mapping_types=mapping_types,
  1008. id=(*id, idx),
  1009. **options,
  1010. )
  1011. )
  1012. return pairs
  1013. elif isinstance(actual, mapping_types) and isinstance(expected, mapping_types):
  1014. actual_keys = set(actual.keys()) # type: ignore[attr-defined]
  1015. expected_keys = set(expected.keys()) # type: ignore[attr-defined]
  1016. if actual_keys != expected_keys:
  1017. missing_keys = expected_keys - actual_keys
  1018. additional_keys = actual_keys - expected_keys
  1019. raise ErrorMeta(
  1020. AssertionError,
  1021. (
  1022. f"The keys of the mappings do not match:\n"
  1023. f"Missing keys in the actual mapping: {sorted(missing_keys)}\n"
  1024. f"Additional keys in the actual mapping: {sorted(additional_keys)}"
  1025. ),
  1026. id=id,
  1027. )
  1028. keys: Collection = actual_keys
  1029. # Since the origination aborts after the first failure, we try to be deterministic
  1030. with contextlib.suppress(Exception):
  1031. keys = sorted(keys)
  1032. pairs = []
  1033. for key in keys:
  1034. pairs.extend(
  1035. originate_pairs(
  1036. actual[key], # type: ignore[index]
  1037. expected[key], # type: ignore[index]
  1038. pair_types=pair_types,
  1039. sequence_types=sequence_types,
  1040. mapping_types=mapping_types,
  1041. id=(*id, key),
  1042. **options,
  1043. )
  1044. )
  1045. return pairs
  1046. else:
  1047. for pair_type in pair_types:
  1048. try:
  1049. # pyrefly: ignore [bad-instantiation]
  1050. return [pair_type(actual, expected, id=id, **options)]
  1051. # Raising an `UnsupportedInputs` during origination indicates that the pair type is not able to handle the
  1052. # inputs. Thus, we try the next pair type.
  1053. except UnsupportedInputs:
  1054. continue
  1055. # Raising an `ErrorMeta` during origination is the orderly way to abort and so we simply re-raise it. This
  1056. # is only in a separate branch, because the one below would also except it.
  1057. except ErrorMeta:
  1058. raise
  1059. # Raising any other exception during origination is unexpected and will give some extra information about
  1060. # what happened. If applicable, the exception should be expected in the future.
  1061. except Exception as error:
  1062. raise RuntimeError(
  1063. f"Originating a {pair_type.__name__}() at item {''.join(str([item]) for item in id)} with\n\n"
  1064. f"{type(actual).__name__}(): {actual}\n\n"
  1065. f"and\n\n"
  1066. f"{type(expected).__name__}(): {expected}\n\n"
  1067. f"resulted in the unexpected exception above. "
  1068. f"If you are a user and see this message during normal operation "
  1069. "please file an issue at https://github.com/pytorch/pytorch/issues. "
  1070. "If you are a developer and working on the comparison functions, "
  1071. "please except the previous error and raise an expressive `ErrorMeta` instead."
  1072. ) from error
  1073. else:
  1074. raise ErrorMeta(
  1075. TypeError,
  1076. f"No comparison pair was able to handle inputs of type {type(actual)} and {type(expected)}.",
  1077. id=id,
  1078. )
  1079. def not_close_error_metas(
  1080. actual: Any,
  1081. expected: Any,
  1082. *,
  1083. pair_types: Sequence[type[Pair]] = (ObjectPair,),
  1084. sequence_types: tuple[type, ...] = (collections.abc.Sequence,),
  1085. mapping_types: tuple[type, ...] = (collections.abc.Mapping,),
  1086. **options: Any,
  1087. ) -> list[ErrorMeta]:
  1088. """Asserts that inputs are equal.
  1089. ``actual`` and ``expected`` can be possibly nested :class:`~collections.abc.Sequence`'s or
  1090. :class:`~collections.abc.Mapping`'s. In this case the comparison happens elementwise by recursing through them.
  1091. Args:
  1092. actual (Any): Actual input.
  1093. expected (Any): Expected input.
  1094. pair_types (Sequence[Type[Pair]]): Sequence of :class:`Pair` types that will be tried to construct with the
  1095. inputs. First successful pair will be used. Defaults to only using :class:`ObjectPair`.
  1096. sequence_types (Tuple[Type, ...]): Optional types treated as sequences that will be checked elementwise.
  1097. mapping_types (Tuple[Type, ...]): Optional types treated as mappings that will be checked elementwise.
  1098. **options (Any): Options passed to each pair during construction.
  1099. """
  1100. # Hide this function from `pytest`'s traceback
  1101. __tracebackhide__ = True
  1102. try:
  1103. pairs = originate_pairs(
  1104. actual,
  1105. expected,
  1106. pair_types=pair_types,
  1107. sequence_types=sequence_types,
  1108. mapping_types=mapping_types,
  1109. **options,
  1110. )
  1111. except ErrorMeta as error_meta:
  1112. # Explicitly raising from None to hide the internal traceback
  1113. raise error_meta.to_error() from None # noqa: RSE102
  1114. error_metas: list[ErrorMeta] = []
  1115. for pair in pairs:
  1116. try:
  1117. pair.compare()
  1118. except ErrorMeta as error_meta:
  1119. error_metas.append(error_meta)
  1120. # Raising any exception besides `ErrorMeta` while comparing is unexpected and will give some extra information
  1121. # about what happened. If applicable, the exception should be expected in the future.
  1122. except Exception as error:
  1123. raise RuntimeError(
  1124. f"Comparing\n\n"
  1125. f"{pair}\n\n"
  1126. f"resulted in the unexpected exception above. "
  1127. f"If you are a user and see this message during normal operation "
  1128. "please file an issue at https://github.com/pytorch/pytorch/issues. "
  1129. "If you are a developer and working on the comparison functions, "
  1130. "please except the previous error and raise an expressive `ErrorMeta` instead."
  1131. ) from error
  1132. # [ErrorMeta Cycles]
  1133. # ErrorMeta objects in this list capture
  1134. # tracebacks that refer to the frame of this function.
  1135. # The local variable `error_metas` refers to the error meta
  1136. # objects, creating a reference cycle. Frames in the traceback
  1137. # would not get freed until cycle collection, leaking cuda memory in tests.
  1138. # We break the cycle by removing the reference to the error_meta objects
  1139. # from this frame as it returns.
  1140. # pyrefly: ignore [bad-assignment]
  1141. error_metas = [error_metas]
  1142. # pyrefly: ignore [bad-return]
  1143. return error_metas.pop()
  1144. def assert_close(
  1145. actual: Any,
  1146. expected: Any,
  1147. *,
  1148. allow_subclasses: bool = True,
  1149. rtol: Optional[float] = None,
  1150. atol: Optional[float] = None,
  1151. equal_nan: bool = False,
  1152. check_device: bool = True,
  1153. check_dtype: bool = True,
  1154. check_layout: bool = True,
  1155. check_stride: bool = False,
  1156. msg: Optional[Union[str, Callable[[str], str]]] = None,
  1157. ):
  1158. r"""Asserts that ``actual`` and ``expected`` are close.
  1159. If ``actual`` and ``expected`` are strided, non-quantized, real-valued, and finite, they are considered close if
  1160. .. math::
  1161. \lvert \text{actual} - \text{expected} \rvert \le \texttt{atol} + \texttt{rtol} \cdot \lvert \text{expected} \rvert
  1162. Non-finite values (``-inf`` and ``inf``) are only considered close if and only if they are equal. ``NaN``'s are
  1163. only considered equal to each other if ``equal_nan`` is ``True``.
  1164. In addition, they are only considered close if they have the same
  1165. - :attr:`~torch.Tensor.device` (if ``check_device`` is ``True``),
  1166. - ``dtype`` (if ``check_dtype`` is ``True``),
  1167. - ``layout`` (if ``check_layout`` is ``True``), and
  1168. - stride (if ``check_stride`` is ``True``).
  1169. If either ``actual`` or ``expected`` is a meta tensor, only the attribute checks will be performed.
  1170. If ``actual`` and ``expected`` are sparse (either having COO, CSR, CSC, BSR, or BSC layout), their strided members are
  1171. checked individually. Indices, namely ``indices`` for COO, ``crow_indices`` and ``col_indices`` for CSR and BSR,
  1172. or ``ccol_indices`` and ``row_indices`` for CSC and BSC layouts, respectively,
  1173. are always checked for equality whereas the values are checked for closeness according to the definition above.
  1174. If ``actual`` and ``expected`` are quantized, they are considered close if they have the same
  1175. :meth:`~torch.Tensor.qscheme` and the result of :meth:`~torch.Tensor.dequantize` is close according to the
  1176. definition above.
  1177. ``actual`` and ``expected`` can be :class:`~torch.Tensor`'s or any tensor-or-scalar-likes from which
  1178. :class:`torch.Tensor`'s can be constructed with :func:`torch.as_tensor`. Except for Python scalars the input types
  1179. have to be directly related. In addition, ``actual`` and ``expected`` can be :class:`~collections.abc.Sequence`'s
  1180. or :class:`~collections.abc.Mapping`'s in which case they are considered close if their structure matches and all
  1181. their elements are considered close according to the above definition.
  1182. .. note::
  1183. Python scalars are an exception to the type relation requirement, because their :func:`type`, i.e.
  1184. :class:`int`, :class:`float`, and :class:`complex`, is equivalent to the ``dtype`` of a tensor-like. Thus,
  1185. Python scalars of different types can be checked, but require ``check_dtype=False``.
  1186. Args:
  1187. actual (Any): Actual input.
  1188. expected (Any): Expected input.
  1189. allow_subclasses (bool): If ``True`` (default) and except for Python scalars, inputs of directly related types
  1190. are allowed. Otherwise type equality is required.
  1191. rtol (Optional[float]): Relative tolerance. If specified ``atol`` must also be specified. If omitted, default
  1192. values based on the :attr:`~torch.Tensor.dtype` are selected with the below table.
  1193. atol (Optional[float]): Absolute tolerance. If specified ``rtol`` must also be specified. If omitted, default
  1194. values based on the :attr:`~torch.Tensor.dtype` are selected with the below table.
  1195. equal_nan (Union[bool, str]): If ``True``, two ``NaN`` values will be considered equal.
  1196. check_device (bool): If ``True`` (default), asserts that corresponding tensors are on the same
  1197. :attr:`~torch.Tensor.device`. If this check is disabled, tensors on different
  1198. :attr:`~torch.Tensor.device`'s are moved to the CPU before being compared.
  1199. check_dtype (bool): If ``True`` (default), asserts that corresponding tensors have the same ``dtype``. If this
  1200. check is disabled, tensors with different ``dtype``'s are promoted to a common ``dtype`` (according to
  1201. :func:`torch.promote_types`) before being compared.
  1202. check_layout (bool): If ``True`` (default), asserts that corresponding tensors have the same ``layout``. If this
  1203. check is disabled, tensors with different ``layout``'s are converted to strided tensors before being
  1204. compared.
  1205. check_stride (bool): If ``True`` and corresponding tensors are strided, asserts that they have the same stride.
  1206. msg (Optional[Union[str, Callable[[str], str]]]): Optional error message to use in case a failure occurs during
  1207. the comparison. Can also passed as callable in which case it will be called with the generated message and
  1208. should return the new message.
  1209. Raises:
  1210. ValueError: If no :class:`torch.Tensor` can be constructed from an input.
  1211. ValueError: If only ``rtol`` or ``atol`` is specified.
  1212. AssertionError: If corresponding inputs are not Python scalars and are not directly related.
  1213. AssertionError: If ``allow_subclasses`` is ``False``, but corresponding inputs are not Python scalars and have
  1214. different types.
  1215. AssertionError: If the inputs are :class:`~collections.abc.Sequence`'s, but their length does not match.
  1216. AssertionError: If the inputs are :class:`~collections.abc.Mapping`'s, but their set of keys do not match.
  1217. AssertionError: If corresponding tensors do not have the same :attr:`~torch.Tensor.shape`.
  1218. AssertionError: If ``check_layout`` is ``True``, but corresponding tensors do not have the same
  1219. :attr:`~torch.Tensor.layout`.
  1220. AssertionError: If only one of corresponding tensors is quantized.
  1221. AssertionError: If corresponding tensors are quantized, but have different :meth:`~torch.Tensor.qscheme`'s.
  1222. AssertionError: If ``check_device`` is ``True``, but corresponding tensors are not on the same
  1223. :attr:`~torch.Tensor.device`.
  1224. AssertionError: If ``check_dtype`` is ``True``, but corresponding tensors do not have the same ``dtype``.
  1225. AssertionError: If ``check_stride`` is ``True``, but corresponding strided tensors do not have the same stride.
  1226. AssertionError: If the values of corresponding tensors are not close according to the definition above.
  1227. The following table displays the default ``rtol`` and ``atol`` for different ``dtype``'s. In case of mismatching
  1228. ``dtype``'s, the maximum of both tolerances is used.
  1229. +---------------------------+------------+----------+
  1230. | ``dtype`` | ``rtol`` | ``atol`` |
  1231. +===========================+============+==========+
  1232. | :attr:`~torch.float16` | ``1e-3`` | ``1e-5`` |
  1233. +---------------------------+------------+----------+
  1234. | :attr:`~torch.bfloat16` | ``1.6e-2`` | ``1e-5`` |
  1235. +---------------------------+------------+----------+
  1236. | :attr:`~torch.float32` | ``1.3e-6`` | ``1e-5`` |
  1237. +---------------------------+------------+----------+
  1238. | :attr:`~torch.float64` | ``1e-7`` | ``1e-7`` |
  1239. +---------------------------+------------+----------+
  1240. | :attr:`~torch.complex32` | ``1e-3`` | ``1e-5`` |
  1241. +---------------------------+------------+----------+
  1242. | :attr:`~torch.complex64` | ``1.3e-6`` | ``1e-5`` |
  1243. +---------------------------+------------+----------+
  1244. | :attr:`~torch.complex128` | ``1e-7`` | ``1e-7`` |
  1245. +---------------------------+------------+----------+
  1246. | :attr:`~torch.quint8` | ``1.3e-6`` | ``1e-5`` |
  1247. +---------------------------+------------+----------+
  1248. | :attr:`~torch.quint2x4` | ``1.3e-6`` | ``1e-5`` |
  1249. +---------------------------+------------+----------+
  1250. | :attr:`~torch.quint4x2` | ``1.3e-6`` | ``1e-5`` |
  1251. +---------------------------+------------+----------+
  1252. | :attr:`~torch.qint8` | ``1.3e-6`` | ``1e-5`` |
  1253. +---------------------------+------------+----------+
  1254. | :attr:`~torch.qint32` | ``1.3e-6`` | ``1e-5`` |
  1255. +---------------------------+------------+----------+
  1256. | other | ``0.0`` | ``0.0`` |
  1257. +---------------------------+------------+----------+
  1258. .. note::
  1259. :func:`~torch.testing.assert_close` is highly configurable with strict default settings. Users are encouraged
  1260. to :func:`~functools.partial` it to fit their use case. For example, if an equality check is needed, one might
  1261. define an ``assert_equal`` that uses zero tolerances for every ``dtype`` by default:
  1262. >>> import functools
  1263. >>> assert_equal = functools.partial(torch.testing.assert_close, rtol=0, atol=0)
  1264. >>> assert_equal(1e-9, 1e-10)
  1265. Traceback (most recent call last):
  1266. ...
  1267. AssertionError: Scalars are not equal!
  1268. <BLANKLINE>
  1269. Expected 1e-10 but got 1e-09.
  1270. Absolute difference: 9.000000000000001e-10
  1271. Relative difference: 9.0
  1272. Examples:
  1273. >>> # tensor to tensor comparison
  1274. >>> expected = torch.tensor([1e0, 1e-1, 1e-2])
  1275. >>> actual = torch.acos(torch.cos(expected))
  1276. >>> torch.testing.assert_close(actual, expected)
  1277. >>> # scalar to scalar comparison
  1278. >>> import math
  1279. >>> expected = math.sqrt(2.0)
  1280. >>> actual = 2.0 / math.sqrt(2.0)
  1281. >>> torch.testing.assert_close(actual, expected)
  1282. >>> # numpy array to numpy array comparison
  1283. >>> import numpy as np
  1284. >>> expected = np.array([1e0, 1e-1, 1e-2])
  1285. >>> actual = np.arccos(np.cos(expected))
  1286. >>> torch.testing.assert_close(actual, expected)
  1287. >>> # sequence to sequence comparison
  1288. >>> import numpy as np
  1289. >>> # The types of the sequences do not have to match. They only have to have the same
  1290. >>> # length and their elements have to match.
  1291. >>> expected = [torch.tensor([1.0]), 2.0, np.array(3.0)]
  1292. >>> actual = tuple(expected)
  1293. >>> torch.testing.assert_close(actual, expected)
  1294. >>> # mapping to mapping comparison
  1295. >>> from collections import OrderedDict
  1296. >>> import numpy as np
  1297. >>> foo = torch.tensor(1.0)
  1298. >>> bar = 2.0
  1299. >>> baz = np.array(3.0)
  1300. >>> # The types and a possible ordering of mappings do not have to match. They only
  1301. >>> # have to have the same set of keys and their elements have to match.
  1302. >>> expected = OrderedDict([("foo", foo), ("bar", bar), ("baz", baz)])
  1303. >>> actual = {"baz": baz, "bar": bar, "foo": foo}
  1304. >>> torch.testing.assert_close(actual, expected)
  1305. >>> expected = torch.tensor([1.0, 2.0, 3.0])
  1306. >>> actual = expected.clone()
  1307. >>> # By default, directly related instances can be compared
  1308. >>> torch.testing.assert_close(torch.nn.Parameter(actual), expected)
  1309. >>> # This check can be made more strict with allow_subclasses=False
  1310. >>> torch.testing.assert_close(
  1311. ... torch.nn.Parameter(actual), expected, allow_subclasses=False
  1312. ... )
  1313. Traceback (most recent call last):
  1314. ...
  1315. TypeError: No comparison pair was able to handle inputs of type
  1316. <class 'torch.nn.parameter.Parameter'> and <class 'torch.Tensor'>.
  1317. >>> # If the inputs are not directly related, they are never considered close
  1318. >>> torch.testing.assert_close(actual.numpy(), expected)
  1319. Traceback (most recent call last):
  1320. ...
  1321. TypeError: No comparison pair was able to handle inputs of type <class 'numpy.ndarray'>
  1322. and <class 'torch.Tensor'>.
  1323. >>> # Exceptions to these rules are Python scalars. They can be checked regardless of
  1324. >>> # their type if check_dtype=False.
  1325. >>> torch.testing.assert_close(1.0, 1, check_dtype=False)
  1326. >>> # NaN != NaN by default.
  1327. >>> expected = torch.tensor(float("Nan"))
  1328. >>> actual = expected.clone()
  1329. >>> torch.testing.assert_close(actual, expected)
  1330. Traceback (most recent call last):
  1331. ...
  1332. AssertionError: Scalars are not close!
  1333. <BLANKLINE>
  1334. Expected nan but got nan.
  1335. Absolute difference: nan (up to 1e-05 allowed)
  1336. Relative difference: nan (up to 1.3e-06 allowed)
  1337. >>> torch.testing.assert_close(actual, expected, equal_nan=True)
  1338. >>> expected = torch.tensor([1.0, 2.0, 3.0])
  1339. >>> actual = torch.tensor([1.0, 4.0, 5.0])
  1340. >>> # The default error message can be overwritten.
  1341. >>> torch.testing.assert_close(
  1342. ... actual, expected, msg="Argh, the tensors are not close!"
  1343. ... )
  1344. Traceback (most recent call last):
  1345. ...
  1346. AssertionError: Argh, the tensors are not close!
  1347. >>> # If msg is a callable, it can be used to augment the generated message with
  1348. >>> # extra information
  1349. >>> torch.testing.assert_close(
  1350. ... actual, expected, msg=lambda msg: f"Header\n\n{msg}\n\nFooter"
  1351. ... )
  1352. Traceback (most recent call last):
  1353. ...
  1354. AssertionError: Header
  1355. <BLANKLINE>
  1356. Tensor-likes are not close!
  1357. <BLANKLINE>
  1358. Mismatched elements: 2 / 3 (66.7%)
  1359. Greatest absolute difference: 2.0 at index (1,) (up to 1e-05 allowed)
  1360. Greatest relative difference: 1.0 at index (1,) (up to 1.3e-06 allowed)
  1361. <BLANKLINE>
  1362. Footer
  1363. """
  1364. # Hide this function from `pytest`'s traceback
  1365. __tracebackhide__ = True
  1366. error_metas = not_close_error_metas(
  1367. actual,
  1368. expected,
  1369. pair_types=(
  1370. NonePair,
  1371. BooleanPair,
  1372. NumberPair,
  1373. TensorLikePair,
  1374. ),
  1375. allow_subclasses=allow_subclasses,
  1376. rtol=rtol,
  1377. atol=atol,
  1378. equal_nan=equal_nan,
  1379. check_device=check_device,
  1380. check_dtype=check_dtype,
  1381. check_layout=check_layout,
  1382. check_stride=check_stride,
  1383. msg=msg,
  1384. )
  1385. if error_metas:
  1386. # TODO: compose all metas into one AssertionError
  1387. raise error_metas[0].to_error(msg)
  1388. @deprecated(
  1389. "`torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. "
  1390. "Please use `torch.testing.assert_close()` instead. "
  1391. "You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.",
  1392. category=FutureWarning,
  1393. )
  1394. def assert_allclose(
  1395. actual: Any,
  1396. expected: Any,
  1397. rtol: Optional[float] = None,
  1398. atol: Optional[float] = None,
  1399. equal_nan: bool = True,
  1400. msg: str = "",
  1401. ) -> None:
  1402. """
  1403. .. warning::
  1404. :func:`torch.testing.assert_allclose` is deprecated since ``1.12`` and will be removed in a future release.
  1405. Please use :func:`torch.testing.assert_close` instead. You can find detailed upgrade instructions
  1406. `here <https://github.com/pytorch/pytorch/issues/61844>`_.
  1407. """
  1408. if not isinstance(actual, torch.Tensor):
  1409. actual = torch.tensor(actual)
  1410. if not isinstance(expected, torch.Tensor):
  1411. expected = torch.tensor(expected, dtype=actual.dtype)
  1412. if rtol is None and atol is None:
  1413. rtol, atol = default_tolerances(
  1414. actual,
  1415. expected,
  1416. dtype_precisions={
  1417. torch.float16: (1e-3, 1e-3),
  1418. torch.float32: (1e-4, 1e-5),
  1419. torch.float64: (1e-5, 1e-8),
  1420. },
  1421. )
  1422. torch.testing.assert_close(
  1423. actual,
  1424. expected,
  1425. rtol=rtol,
  1426. atol=atol,
  1427. equal_nan=equal_nan,
  1428. check_device=True,
  1429. check_dtype=False,
  1430. check_stride=False,
  1431. msg=msg or None,
  1432. )