chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,682 @@
|
||||
# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import importlib.abc
|
||||
import importlib.util
|
||||
import inspect
|
||||
import pkgutil
|
||||
import sys
|
||||
import types
|
||||
import warnings
|
||||
from contextlib import contextmanager
|
||||
from functools import cache
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Generator, Iterable
|
||||
from typing import TypeAlias
|
||||
|
||||
_ScopeType: TypeAlias = str | Iterable[str] | None
|
||||
|
||||
|
||||
def warning_about_fake_interface(name: str):
|
||||
warnings.warn(
|
||||
f"The interface '{name}' is a fake implementation for torch compatibility. "
|
||||
"It does not have the actual functionality of PyTorch. "
|
||||
"Please refer to the PaddlePaddle documentation for equivalent functionality.",
|
||||
category=UserWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
|
||||
def create_fake_class(name, attrs: dict[str, Any]):
|
||||
"""Create a fake class with the given name and attributes."""
|
||||
new_fn = lambda *args, **kwargs: warning_about_fake_interface(name)
|
||||
attrs["__init__"] = new_fn
|
||||
return type(name, (), attrs)
|
||||
|
||||
|
||||
def create_fake_function(name):
|
||||
"""Create a fake function with the given name and implementation."""
|
||||
fn = lambda *args, **kwargs: warning_about_fake_interface(name)
|
||||
fn.__name__ = name
|
||||
return fn
|
||||
|
||||
|
||||
class OverriddenAttribute:
|
||||
def get_value(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class LazyImportOverriddenAttribute(OverriddenAttribute):
|
||||
def __init__(self, full_name: str):
|
||||
self._full_name = full_name
|
||||
|
||||
def get_value(self):
|
||||
parts = self._full_name.split(".")
|
||||
root_module = importlib.import_module(parts[0])
|
||||
result = root_module
|
||||
for part in parts[1:]:
|
||||
result = getattr(result, part)
|
||||
return result
|
||||
|
||||
|
||||
class RawOverriddenAttribute(OverriddenAttribute):
|
||||
def __init__(self, value: Any):
|
||||
self._value = value
|
||||
|
||||
def get_value(self):
|
||||
return self._value
|
||||
|
||||
|
||||
class ProxyModule(types.ModuleType):
|
||||
def __init__(
|
||||
self,
|
||||
original_module: types.ModuleType,
|
||||
proxy_name: str,
|
||||
overrides: dict[str, OverriddenAttribute],
|
||||
):
|
||||
super().__init__(proxy_name)
|
||||
self._original_module = original_module
|
||||
self._proxy_name = proxy_name
|
||||
self._overrides = overrides
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
if name in self._overrides:
|
||||
return self._overrides[name].get_value()
|
||||
return getattr(self._original_module, name)
|
||||
|
||||
|
||||
class CallableProxyModule(ProxyModule):
|
||||
"""
|
||||
Preserve callability for modules whose type defines ``__call__``.
|
||||
|
||||
``callable(obj)`` does not consult ``obj.__getattr__("__call__")``. It checks the
|
||||
type-level call slot instead, so callable modules need a dedicated proxy subtype.
|
||||
"""
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
return self._original_module(*args, **kwargs)
|
||||
|
||||
|
||||
def _create_proxy_module(
|
||||
original_module: types.ModuleType,
|
||||
proxy_name: str,
|
||||
overrides: dict[str, OverriddenAttribute],
|
||||
) -> ProxyModule:
|
||||
"""Wrap callable modules with a callable proxy and plain modules with ProxyModule."""
|
||||
if callable(original_module):
|
||||
return CallableProxyModule(original_module, proxy_name, overrides)
|
||||
return ProxyModule(original_module, proxy_name, overrides)
|
||||
|
||||
|
||||
GLOBAL_OVERRIDES: dict[str, OverriddenAttribute] = {
|
||||
"torch.relu": LazyImportOverriddenAttribute("paddle.nn.functional.relu"),
|
||||
"torch.TorchVersion": LazyImportOverriddenAttribute(
|
||||
"paddle.paddle_version.PaddleVersion"
|
||||
),
|
||||
}
|
||||
|
||||
TORCH_PROXY_BLOCKED_MODULES = {
|
||||
"tvm_ffi",
|
||||
}
|
||||
|
||||
MAGIC_DISABLED_MODULE_ATTR: str = "__disable_torch_proxy__"
|
||||
MAGIC_ENABLED_MODULE_ATTR: str = "__enable_torch_proxy__"
|
||||
|
||||
|
||||
def _extend_torch_proxy_overrides(
|
||||
overrides: dict[str, OverriddenAttribute],
|
||||
) -> None:
|
||||
GLOBAL_OVERRIDES.update(overrides)
|
||||
|
||||
|
||||
@cache
|
||||
def _register_compat_override():
|
||||
import paddle.compat
|
||||
|
||||
PADDLE_PREFIX = "paddle.compat"
|
||||
TORCH_PREFIX = "torch"
|
||||
PUBLIC_ATTR_DECLARATION = "__all__"
|
||||
|
||||
compat_overrides = {}
|
||||
for module_info in pkgutil.walk_packages(
|
||||
paddle.compat.__path__,
|
||||
paddle.compat.__name__ + ".",
|
||||
):
|
||||
module = importlib.import_module(module_info.name)
|
||||
if hasattr(module, PUBLIC_ATTR_DECLARATION):
|
||||
public_attrs = getattr(module, PUBLIC_ATTR_DECLARATION)
|
||||
torch_module_name = module_info.name.replace(
|
||||
PADDLE_PREFIX, TORCH_PREFIX, 1
|
||||
)
|
||||
for attr_name in public_attrs:
|
||||
if attr_name.startswith("_"):
|
||||
continue
|
||||
paddle_attr = getattr(module, attr_name)
|
||||
torch_attr_name = f"{torch_module_name}.{attr_name}"
|
||||
compat_overrides[torch_attr_name] = RawOverriddenAttribute(
|
||||
paddle_attr
|
||||
)
|
||||
_extend_torch_proxy_overrides(compat_overrides)
|
||||
|
||||
|
||||
def _is_specific_module_or_its_submodule(name: str, module: str) -> bool:
|
||||
return name == module or name.startswith(f"{module}.")
|
||||
|
||||
|
||||
def _is_torch_module(name: str) -> bool:
|
||||
return _is_specific_module_or_its_submodule(name, "torch")
|
||||
|
||||
|
||||
def _is_torch_proxy_local_enabled_module(name: str, scope: set[str]) -> bool:
|
||||
for enabled_module in scope:
|
||||
if _is_specific_module_or_its_submodule(name, enabled_module):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_torch_proxy_blocked_module(name: str) -> bool:
|
||||
for blocked_module in TORCH_PROXY_BLOCKED_MODULES:
|
||||
if _is_specific_module_or_its_submodule(name, blocked_module):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_called_by_module_with_specific_dunder_attr(dunder_attr: str) -> bool:
|
||||
frame = inspect.currentframe()
|
||||
while frame is not None:
|
||||
if frame.f_globals.get(dunder_attr):
|
||||
return True
|
||||
frame = frame.f_back
|
||||
return False
|
||||
|
||||
|
||||
def _is_called_by_torch_proxy_blocked_module():
|
||||
return _is_called_by_module_with_specific_dunder_attr(
|
||||
MAGIC_DISABLED_MODULE_ATTR
|
||||
)
|
||||
|
||||
|
||||
def _is_called_by_torch_proxy_local_enabled_module():
|
||||
return _is_called_by_module_with_specific_dunder_attr(
|
||||
MAGIC_ENABLED_MODULE_ATTR
|
||||
)
|
||||
|
||||
|
||||
class TorchProxyMetaFinder:
|
||||
"""
|
||||
PyTorch compatibility layer for PaddlePaddle.
|
||||
|
||||
This class provides a way to `import torch` but actually loads PaddlePaddle.
|
||||
|
||||
Inspired by the setuptools _distutils_hack.
|
||||
"""
|
||||
|
||||
_local_enabled_scope: set[str]
|
||||
_globally_enabled: bool
|
||||
|
||||
def __init__(self, scope: set[str] | None = None):
|
||||
self._set_scope(scope)
|
||||
|
||||
def _set_scope(self, scope: set[str] | None):
|
||||
self._local_enabled_scope = scope or set()
|
||||
self._globally_enabled = scope is None
|
||||
|
||||
def find_spec(self, fullname, path, target=None):
|
||||
if _is_torch_proxy_blocked_module(fullname):
|
||||
return self._find_spec_for_torch_proxy_blocked_module(fullname)
|
||||
|
||||
if _is_torch_proxy_local_enabled_module(
|
||||
fullname, self._local_enabled_scope
|
||||
):
|
||||
return self._find_spec_for_torch_proxy_local_enabled_module(
|
||||
fullname
|
||||
)
|
||||
|
||||
if not _is_torch_module(fullname):
|
||||
return None
|
||||
|
||||
if _is_called_by_torch_proxy_blocked_module():
|
||||
if fullname in TORCH_MODULES_CACHE:
|
||||
return self._find_spec_for_cached_torch_module(fullname)
|
||||
return None
|
||||
|
||||
if (
|
||||
not self._globally_enabled
|
||||
and not _is_called_by_torch_proxy_local_enabled_module()
|
||||
):
|
||||
if fullname in TORCH_MODULES_CACHE:
|
||||
return self._find_spec_for_cached_torch_module(fullname)
|
||||
return None
|
||||
|
||||
return self._find_spec_for_torch_module(fullname)
|
||||
|
||||
def _find_spec_for_specific_module(
|
||||
self,
|
||||
fullname: str,
|
||||
enable_proxy_when_exec_module: bool,
|
||||
patched_dunder_attr: str,
|
||||
):
|
||||
# Return a special loader that imports the blocked module without torch proxy
|
||||
with use_compat_guard(enable=False):
|
||||
spec = importlib.util.find_spec(fullname)
|
||||
if spec is None:
|
||||
return None
|
||||
original_loader = spec.loader
|
||||
if original_loader is None:
|
||||
return None
|
||||
|
||||
class SpecificModuleLoader(importlib.abc.Loader):
|
||||
def create_module(self, spec):
|
||||
mod = original_loader.create_module(spec)
|
||||
if mod is None:
|
||||
# If original loader returns None, create default module
|
||||
# and ensure it has necessary attributes from spec
|
||||
mod = types.ModuleType(spec.name)
|
||||
mod.__spec__ = spec
|
||||
mod.__loader__ = self
|
||||
if spec.origin is not None:
|
||||
mod.__file__ = spec.origin
|
||||
if spec.submodule_search_locations is not None:
|
||||
mod.__path__ = list(spec.submodule_search_locations)
|
||||
return mod
|
||||
|
||||
def exec_module(self, module):
|
||||
# Import the real module with torch proxy disabled
|
||||
with use_compat_guard(
|
||||
enable=enable_proxy_when_exec_module, silent=True
|
||||
):
|
||||
original_loader.exec_module(module)
|
||||
# Mark module as torch proxy disabled/local enabled
|
||||
module.__dict__[patched_dunder_attr] = True
|
||||
|
||||
spec.loader = SpecificModuleLoader()
|
||||
return spec
|
||||
|
||||
def _find_spec_for_torch_proxy_local_enabled_module(self, fullname: str):
|
||||
return self._find_spec_for_specific_module(
|
||||
fullname,
|
||||
enable_proxy_when_exec_module=True,
|
||||
patched_dunder_attr=MAGIC_ENABLED_MODULE_ATTR,
|
||||
)
|
||||
|
||||
def _find_spec_for_torch_proxy_blocked_module(self, fullname: str):
|
||||
return self._find_spec_for_specific_module(
|
||||
fullname,
|
||||
enable_proxy_when_exec_module=False,
|
||||
patched_dunder_attr=MAGIC_DISABLED_MODULE_ATTR,
|
||||
)
|
||||
|
||||
def _find_spec_for_cached_torch_module(self, fullname: str):
|
||||
module = TORCH_MODULES_CACHE[fullname]
|
||||
|
||||
# Return cached module before enable proxy
|
||||
class CachedTorchModuleLoader(importlib.abc.Loader):
|
||||
def create_module(self, spec):
|
||||
return module
|
||||
|
||||
def exec_module(self, module):
|
||||
pass
|
||||
|
||||
# Always treat cached modules as packages to allow submodules to be loaded.
|
||||
# This is necessary because some modules (e.g. torch._C) are not packages
|
||||
# but have submodules (e.g. torch._C._dynamo) attached to them.
|
||||
spec = importlib.util.spec_from_loader(
|
||||
fullname,
|
||||
CachedTorchModuleLoader(),
|
||||
origin=getattr(module, "__file__", None),
|
||||
is_package=True,
|
||||
)
|
||||
spec.submodule_search_locations = list(getattr(module, "__path__", []))
|
||||
return spec
|
||||
|
||||
def _find_spec_for_torch_module(self, fullname: str):
|
||||
# Map the requested torch fullname to the corresponding paddle fullname.
|
||||
module_name = fullname.replace("torch", "paddle", 1)
|
||||
source_module = importlib.import_module(module_name)
|
||||
overrides = {
|
||||
k.removeprefix(f"{fullname}."): v
|
||||
for k, v in GLOBAL_OVERRIDES.items()
|
||||
if k.startswith(f"{fullname}.")
|
||||
}
|
||||
|
||||
is_pkg = hasattr(source_module, "__path__")
|
||||
|
||||
class TorchProxyLoader(importlib.abc.Loader):
|
||||
def __init__(self, source, target_name):
|
||||
self._source = source
|
||||
self._target_name = target_name
|
||||
|
||||
def create_module(self, spec):
|
||||
# Create a new module object that will act as the "torch..." module.
|
||||
mod = _create_proxy_module(
|
||||
self._source, self._target_name, overrides
|
||||
)
|
||||
# Preserve file/path information for tooling/debugging.
|
||||
mod.__file__ = getattr(self._source, "__file__", None)
|
||||
if is_pkg:
|
||||
# package must expose __path__ so import machinery can find submodules
|
||||
mod.__path__ = list(getattr(self._source, "__path__", []))
|
||||
mod.__package__ = self._target_name
|
||||
else:
|
||||
mod.__package__ = self._target_name.rpartition('.')[0]
|
||||
return mod
|
||||
|
||||
def exec_module(self, module):
|
||||
# Populate the new module with attributes from the source paddle module.
|
||||
# Skip a few special attributes that should reflect the new module name.
|
||||
for k, v in self._source.__dict__.items():
|
||||
if k in ("__name__", "__package__", "__path__", "__spec__"):
|
||||
continue
|
||||
if k in overrides:
|
||||
continue
|
||||
if isinstance(v, types.ModuleType):
|
||||
v = _create_proxy_module(
|
||||
v,
|
||||
f"{self._target_name}.{k}",
|
||||
{
|
||||
kk.removeprefix(f"{k}."): vv
|
||||
for kk, vv in overrides.items()
|
||||
if kk.startswith(f"{k}.")
|
||||
},
|
||||
)
|
||||
module.__dict__[k] = v
|
||||
|
||||
# Use fullname for the spec name and mark as package when appropriate so that
|
||||
# statements like `import torch.nn.functional` work correctly.
|
||||
return importlib.util.spec_from_loader(
|
||||
fullname,
|
||||
TorchProxyLoader(source_module, fullname),
|
||||
is_package=is_pkg,
|
||||
origin=getattr(source_module, "__file__", None),
|
||||
)
|
||||
|
||||
|
||||
TORCH_PROXY_FINDER = TorchProxyMetaFinder()
|
||||
TORCH_MODULES_CACHE: dict[str, types.ModuleType] = {}
|
||||
|
||||
|
||||
def _clear_torch_proxy_modules():
|
||||
for name, module in list(sys.modules.items()):
|
||||
if _is_torch_module(name) and isinstance(module, ProxyModule):
|
||||
del sys.modules[name]
|
||||
|
||||
|
||||
def _swap_torch_modules_to_cache():
|
||||
for name, module in list(sys.modules.items()):
|
||||
if _is_torch_module(name):
|
||||
if not isinstance(module, ProxyModule):
|
||||
TORCH_MODULES_CACHE[name] = sys.modules[name]
|
||||
del sys.modules[name]
|
||||
|
||||
|
||||
def _copy_torch_modules_from_cache():
|
||||
for name in list(TORCH_MODULES_CACHE):
|
||||
assert _is_torch_module(name), f"`{name}` is not a PyTorch module"
|
||||
sys.modules[name] = TORCH_MODULES_CACHE[name]
|
||||
|
||||
|
||||
def _modify_scope_of_torch_proxy(
|
||||
scope: set[str] | None,
|
||||
*,
|
||||
silent: bool = False,
|
||||
) -> None:
|
||||
def _warn_or_not(msg: str):
|
||||
if silent:
|
||||
return
|
||||
warnings.warn(msg)
|
||||
|
||||
if TORCH_PROXY_FINDER not in sys.meta_path:
|
||||
TORCH_PROXY_FINDER._set_scope(scope)
|
||||
return
|
||||
|
||||
if TORCH_PROXY_FINDER._globally_enabled:
|
||||
if scope is not None:
|
||||
_warn_or_not(
|
||||
"PyTorch already enabled globally, scope modification ignored."
|
||||
)
|
||||
TORCH_PROXY_FINDER._set_scope(scope)
|
||||
return
|
||||
if scope is None:
|
||||
_warn_or_not(
|
||||
"Enabling PyTorch compat globally, previous scope will be ignored."
|
||||
)
|
||||
TORCH_PROXY_FINDER._globally_enabled = True
|
||||
return
|
||||
if scope != TORCH_PROXY_FINDER._local_enabled_scope:
|
||||
_warn_or_not(
|
||||
f"Extending PyTorch compat scope, previous scope: {TORCH_PROXY_FINDER._local_enabled_scope}, new scope: {scope}."
|
||||
)
|
||||
TORCH_PROXY_FINDER._local_enabled_scope |= scope
|
||||
|
||||
|
||||
def _parse_scope(scope: str | Iterable[str] | None) -> set[str] | None:
|
||||
if scope is None:
|
||||
return None
|
||||
if isinstance(scope, str):
|
||||
return {scope}
|
||||
return set(scope)
|
||||
|
||||
|
||||
def enable_compat(
|
||||
*,
|
||||
scope: _ScopeType = None,
|
||||
blocked_modules: _ScopeType = None,
|
||||
backend: Literal["torch"] = "torch",
|
||||
silent: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Enable the PyTorch compat by adding the TorchProxyMetaFinder to sys.meta_path.
|
||||
This allows importing 'torch' modules that are actually proxies to PaddlePaddle.
|
||||
|
||||
Args:
|
||||
scope (str or Iterable[str], optional): Specific module or modules to enable
|
||||
PyTorch compat for. If None, enables PyTorch compat globally. Defaults to None.
|
||||
blocked_modules (str or Iterable[str], optional): Specific module or modules to
|
||||
exclude from PyTorch compat. Defaults to None.
|
||||
backend (str, optional): The backend to enable compat for. Currently only
|
||||
"torch" is supported. Defaults to "torch".
|
||||
silent (bool, optional): If True, suppresses warnings about scope changes.
|
||||
Defaults to False.
|
||||
|
||||
Example:
|
||||
.. code-block:: pycon
|
||||
:name: enable-compat-in-global-scope
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.enable_compat() # Enable torch compat globally
|
||||
>>> import torch # type: ignore[import-not-found] # This will import paddle as torch
|
||||
>>> assert torch.sin is paddle.sin
|
||||
>>> paddle.disable_compat() # Disable torch compat
|
||||
|
||||
.. code-block:: pycon
|
||||
:name: enable-compat-in-specific-scope
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.enable_compat(scope={"triton"}) # Enable torch compat for 'triton' module only
|
||||
>>> import triton # type: ignore[import-untyped] # All `import torch` inside `triton` will proxy to paddle
|
||||
>>> try:
|
||||
... import torch # type: ignore[import-not-found] # This will raise ModuleNotFoundError
|
||||
... except ModuleNotFoundError:
|
||||
... print("PyTorch compat is not enabled globally.")
|
||||
>>> paddle.disable_compat() # Disable torch compat
|
||||
"""
|
||||
assert backend == "torch", f"Unsupported backend: {backend}"
|
||||
blocked_modules = _parse_scope(blocked_modules)
|
||||
if blocked_modules is not None:
|
||||
extend_torch_proxy_blocked_modules(blocked_modules)
|
||||
scope = _parse_scope(scope)
|
||||
_register_compat_override()
|
||||
_swap_torch_modules_to_cache()
|
||||
_modify_scope_of_torch_proxy(scope, silent=silent)
|
||||
sys.meta_path.insert(0, TORCH_PROXY_FINDER)
|
||||
|
||||
|
||||
def disable_compat() -> None:
|
||||
"""
|
||||
Disable the PyTorch proxy by removing the TorchProxyMetaFinder from sys.meta_path.
|
||||
This prevents 'torch' imports from being proxied to PaddlePaddle.
|
||||
|
||||
Example:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.enable_compat() # Enable torch compat globally
|
||||
>>> import torch # type: ignore[import-not-found] # This will import paddle as torch
|
||||
>>> assert torch.sin is paddle.sin
|
||||
>>> paddle.disable_compat() # Disable torch compat
|
||||
>>> try:
|
||||
... import torch # This will raise ModuleNotFoundError
|
||||
... except ModuleNotFoundError:
|
||||
... print("PyTorch compat is disabled.")
|
||||
"""
|
||||
if TORCH_PROXY_FINDER in sys.meta_path:
|
||||
sys.meta_path.remove(TORCH_PROXY_FINDER)
|
||||
_clear_torch_proxy_modules()
|
||||
_copy_torch_modules_from_cache()
|
||||
return
|
||||
warnings.warn("torch compat is not installed.")
|
||||
|
||||
|
||||
@contextmanager
|
||||
def use_compat_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.use_compat_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.use_compat_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_compat(scope=scope, silent=silent)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
TORCH_PROXY_FINDER._local_enabled_scope = (
|
||||
original_local_enabled_scope
|
||||
)
|
||||
TORCH_PROXY_FINDER._globally_enabled = original_globally_enabled
|
||||
disable_compat()
|
||||
else:
|
||||
disable_compat()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
enable_compat(scope=None, silent=True)
|
||||
TORCH_PROXY_FINDER._local_enabled_scope = (
|
||||
original_local_enabled_scope
|
||||
)
|
||||
TORCH_PROXY_FINDER._globally_enabled = original_globally_enabled
|
||||
|
||||
|
||||
def extend_torch_proxy_blocked_modules(modules: Iterable[str]) -> None:
|
||||
"""Add modules to the PyTorch proxy blocked list.
|
||||
|
||||
Modules in the blocked list will not use PyTorch compat when imported,
|
||||
and their functions will not trigger PyTorch compat when called.
|
||||
|
||||
By default, some modules are already in the blocked list, such as 'tvm_ffi'.
|
||||
|
||||
Args:
|
||||
modules(Iterable[str]): An iterable of module names to block from PyTorch compat.
|
||||
|
||||
Example:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> paddle.enable_compat() # Enable torch compat globally
|
||||
>>> # Add 'my_custom_module' to the blocked list
|
||||
>>> paddle.compat.extend_torch_proxy_blocked_modules(['my_custom_module'])
|
||||
>>> # doctest: +SKIP('my_custom_module is not available')
|
||||
>>> import my_custom_module # type: ignore[import-not-found] # This import will not use torch compat
|
||||
"""
|
||||
TORCH_PROXY_BLOCKED_MODULES.update(modules)
|
||||
|
||||
|
||||
def paddle_triton_fun():
|
||||
"""
|
||||
Enable the triton support and return triton module.
|
||||
Args: None.
|
||||
Returns: triton module
|
||||
|
||||
Example:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:GPU)
|
||||
>>> from paddle.compat import paddle_triton_fun
|
||||
>>> triton = paddle_triton_fun()
|
||||
>>> import triton.language as tl
|
||||
|
||||
>>> @triton.jit
|
||||
>>> def add_kernel(X, Y, Z, N, BLOCK: tl.constexpr):
|
||||
... pid = tl.program_id(0)
|
||||
... offs = pid * BLOCK + tl.arange(0, BLOCK)
|
||||
... mask = offs < N
|
||||
... x = tl.load(X + offs, mask=mask)
|
||||
... y = tl.load(Y + offs, mask=mask)
|
||||
... tl.store(Z + offs, x + y, mask=mask)
|
||||
"""
|
||||
enable_compat(scope={"triton"})
|
||||
import triton
|
||||
|
||||
return triton
|
||||
Reference in New Issue
Block a user