Function use_torch_proxy_guard

python/paddle/compat/proxy.py:530–602

Context manager to temporarily enable or disable the PyTorch compat. When `enable` is True (default), the PyTorch compat is enabled for the duration of the context and restored to its previous state afterwards. When `enable` is False, the PyTorch compat is disabled for the duration

(
    *,
    enable: bool = True,
    scope: _ScopeType = None,
    silent: bool = False,
)

@contextmanager
def use_torch_proxy_guard(
    *,
    enable: bool = True,
    scope: _ScopeType = None,
    silent: bool = False,
) -> Generator[None, None, None]:
    """
    Context manager to temporarily enable or disable the PyTorch compat.

    When `enable` is True (default), the PyTorch compat is enabled for the duration
    of the context and restored to its previous state afterwards. When `enable`
    is False, the PyTorch compat is disabled for the duration of the context and
    restored afterwards.

    Args:
        enable (bool, optional): Whether to enable or disable the PyTorch compat
            within the context. Defaults to True.
        scope (str or Iterable[str], optional): Specific module or modules to enable
            PyTorch compat for. If None, uses the global scope. Defaults to None.
        silent (bool, optional): If True, suppresses warnings about scope changes.
            Defaults to False.

    Example:
        .. code-block:: pycon

            >>> import paddle

            >>> with paddle.compat.use_torch_proxy_guard():
            ...     # code that requires the Torch compat to be enabled
            ...     import torch  # type: ignore[import-not-found]
            ...
            ...     assert torch.sin is paddle.sin
            ...     # Temporarily disable the Torch compat
            ...     with paddle.compat.use_torch_proxy_guard(enable=False):
            ...         try:
            ...             import torch
            ...         except ModuleNotFoundError:
            ...             print("Torch compat is disabled within this block.")
            ...     # Torch compat is re-enabled here
            ...     import torch
            ...
            ...     assert torch.sin is paddle.sin
    """
    scope = _parse_scope(scope)
    already_has_torch_proxy = TORCH_PROXY_FINDER in sys.meta_path
    original_local_enabled_scope = set(TORCH_PROXY_FINDER._local_enabled_scope)
    original_globally_enabled = TORCH_PROXY_FINDER._globally_enabled
    if enable == already_has_torch_proxy and (
        (original_globally_enabled and scope is None)
        or (original_local_enabled_scope == (scope or set()))
    ):
        yield
        return
    if enable:
        enable_torch_proxy(scope=scope, silent=silent)
        try:
            yield
        finally: