xref: /aosp_15_r20/external/pytorch/torch/fx/_compatibility.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1from typing import Any, Dict, Callable, TypeVar
2import textwrap
3
4_BACK_COMPAT_OBJECTS : Dict[Any, None] = {}
5_MARKED_WITH_COMPATIBILITY : Dict[Any, None] = {}
6
7_T = TypeVar("_T")
8
9def compatibility(is_backward_compatible: bool) -> Callable[[_T], _T]:
10    if is_backward_compatible:
11
12        def mark_back_compat(fn: _T) -> _T:
13            docstring = textwrap.dedent(getattr(fn, '__doc__', None) or '')
14            docstring += """
15.. note::
16    Backwards-compatibility for this API is guaranteed.
17"""
18            fn.__doc__ = docstring
19            _BACK_COMPAT_OBJECTS.setdefault(fn)
20            _MARKED_WITH_COMPATIBILITY.setdefault(fn)
21            return fn
22
23        return mark_back_compat
24    else:
25
26        def mark_not_back_compat(fn: _T) -> _T:
27            docstring = textwrap.dedent(getattr(fn, '__doc__', None) or '')
28            docstring += """
29.. warning::
30    This API is experimental and is *NOT* backward-compatible.
31"""
32            fn.__doc__ = docstring
33            _MARKED_WITH_COMPATIBILITY.setdefault(fn)
34            return fn
35
36        return mark_not_back_compat
37