wehub-resource-sync
683 lines
24 KiB
Python
683 lines
24 KiB
Python
# 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
|