[API Compatibility] enhance paddle.enable_compat -part

python/paddle/compat/api_dispatch.py

@@ -0,0 +1,189 @@
+# Copyright (c) 2026 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.
+
+"""``enable_compat(level=2)`` API dispatch: route ``paddle.*`` (and
+``paddle.Tensor`` methods) to the torch-aligned ``paddle.compat.*`` for
+external callers while paddle-internal callers keep native semantics.
+
+The enable/disable/level lifecycle lives in ``paddle.compat.proxy``; this
+module only installs/removes the dispatchers and holds the dispatch state.
+"""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+import pkgutil
+import sys
+from functools import wraps
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import types
+    from collections.abc import Generator
+
+
+# (live module, attr name) -> original value for every patched ``paddle.*``
+# symbol, so ``disable_compat()`` can restore the paddle namespace.
+_PADDLE_NAMESPACE_SAVED: dict[tuple[types.ModuleType, str], Any] = {}
+
+
+def _caller_is_paddle_internal() -> bool:
+    """True when the caller of the dispatched ``paddle.*`` API is a ``paddle``
+    module. Paddle's own internals call these APIs with native kwargs/defaults, so
+    they (and the compat impls) get native; only external callers see compat."""
+    # _getframe(1) = our caller (the dispatch site); its f_back = who called paddle.X.
+    frame = sys._getframe(1).f_back
+    if frame is None:
+        return False
+    name = frame.f_globals.get("__name__") or ""
+    return name == "paddle" or name.startswith("paddle.")
+
+
+def dispatch_function(compat_fn: Any) -> Any:
+    """Wrap a native ``paddle`` callable to route external callers to
+    ``compat_fn`` while compat is enabled; paddle-internal callers and the
+    disabled state get the native callable. Installed only under
+    ``enable_compat(level=2)``; ``disable_compat`` restores the originals,
+    so the default hot path is untouched."""
+
+    def decorator(native_fn: Any) -> Any:
+        @wraps(native_fn)
+        def dispatcher(*args: Any, **kwargs: Any) -> Any:
+            if (
+                len(_PADDLE_NAMESPACE_SAVED) > 0
+                and not _caller_is_paddle_internal()
+            ):
+                return compat_fn(*args, **kwargs)
+            return native_fn(*args, **kwargs)
+
+        dispatcher.__compat_fn__ = compat_fn
+        dispatcher.__native_fn__ = native_fn
+        dispatcher.__signature__ = inspect.signature(compat_fn)
+        return dispatcher
+
+    return decorator
+
+
+def _iter_compat_modules() -> Generator[types.ModuleType, None, None]:
+    """Yield ``paddle.compat`` modules that declare ``__all__``.
+
+    ``pkgutil.walk_packages`` skips the starting package, so the root
+    ``paddle.compat`` (holding the top-level functions) is yielded explicitly.
+    """
+    import paddle.compat
+
+    if hasattr(paddle.compat, "__all__"):
+        yield paddle.compat
+    for module_info in pkgutil.walk_packages(
+        paddle.compat.__path__,
+        paddle.compat.__name__ + ".",
+    ):
+        compat_module = importlib.import_module(module_info.name)
+        if not hasattr(compat_module, "__all__"):
+            continue
+        yield compat_module
+
+
+def dispatch_class(native_cls: type, compat_cls: type) -> type:
+    """Create a class proxy that selects compat only for external callers."""
+
+    class _CompatAwareMeta(type(compat_cls)):
+        def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+            if cls is proxy:
+                if (
+                    len(_PADDLE_NAMESPACE_SAVED) > 0
+                    and not _caller_is_paddle_internal()
+                ):
+                    return compat_cls(*args, **kwargs)
+                return native_cls(*args, **kwargs)
+            return super().__call__(*args, **kwargs)
+
+    proxy = _CompatAwareMeta(
+        native_cls.__name__,
+        (compat_cls,),
+        {
+            "__module__": native_cls.__module__,
+            "__compat_cls__": compat_cls,
+            "__native_cls__": native_cls,
+        },
+    )
+    proxy.__signature__ = inspect.signature(compat_cls)
+    return proxy
+
+
+def _patch_tensor_methods() -> None:
+    """Route ``paddle.Tensor.<m>`` to the compat function for the root compat APIs
+    that torch also exposes as Tensor methods (max/min/sort/split/unique/...), so
+    ``x.max(dim=1)`` works torch-style for external callers (native for internal).
+    The dispatcher is patched directly like any paddle Tensor method: the
+    descriptor protocol forwards the tensor as the first positional argument,
+    which is exactly the compat function's ``input`` parameter.
+    """
+    import paddle
+    import paddle.compat as compat_root
+
+    for attr_name in getattr(compat_root, "__all__", ()):
+        native_method = getattr(paddle.Tensor, attr_name, None)
+        if native_method is None:
+            continue
+        compat_fn = getattr(compat_root, attr_name)
+        _PADDLE_NAMESPACE_SAVED[(paddle.Tensor, attr_name)] = native_method
+        setattr(
+            paddle.Tensor,
+            attr_name,
+            dispatch_function(compat_fn)(native_method),
+        )
+
+
+def _apply_paddle_namespace_aliases() -> None:
+    """Install caller-aware dispatchers/proxies for every public ``paddle.compat.*``
+    symbol that has a ``paddle.*`` counterpart, plus the Tensor methods."""
+    if _PADDLE_NAMESPACE_SAVED:
+        return
+
+    for compat_module in _iter_compat_modules():
+        target_name = compat_module.__name__.replace(
+            "paddle.compat", "paddle", 1
+        )
+        try:
+            target_module = importlib.import_module(target_name)
+        except ModuleNotFoundError:
+            continue
+        for attr_name in compat_module.__all__:
+            compat_attr = getattr(compat_module, attr_name)
+            current = getattr(target_module, attr_name, None)
+            if current is None or current is compat_attr:
+                continue
+            _PADDLE_NAMESPACE_SAVED[(target_module, attr_name)] = current
+            if isinstance(compat_attr, type):
+                setattr(
+                    target_module,
+                    attr_name,
+                    dispatch_class(current, compat_attr),
+                )
+            else:
+                setattr(
+                    target_module,
+                    attr_name,
+                    dispatch_function(compat_attr)(current),
+                )
+    _patch_tensor_methods()
+
+
+def _restore_paddle_namespace_aliases() -> None:
+    """Undo :func:`_apply_paddle_namespace_aliases`, restoring the paddle namespace."""
+    for (target_module, attr_name), original in _PADDLE_NAMESPACE_SAVED.items():
+        setattr(target_module, attr_name, original)
+    _PADDLE_NAMESPACE_SAVED.clear()

python/paddle/compat/distributions/categorical.py

@@ -20,8 +20,10 @@
 from paddle.distribution import constraint, distribution
 from paddle.tensor import multinomial
 
+from ..utils import _CompatClassMeta
 
-class Categorical(distribution.Distribution):
+
+class Categorical(distribution.Distribution, metaclass=_CompatClassMeta):
     arg_constraints = {
         "probs": constraint.simplex,
         "logits": constraint.real_vector,
@@ -61,7 +63,9 @@ def __init__(
         batch_shape = (
             tuple(self._param.shape[:-1]) if self._param.dim() > 1 else ()
         )
-        super().__init__(batch_shape, validate_args=validate_args)
+        distribution.Distribution.__init__(
+            self, batch_shape, validate_args=validate_args
+        )
 
     def expand(self, batch_shape, _instance=None):
         new = (
@@ -81,7 +85,9 @@ def expand(self, batch_shape, _instance=None):
         )
         new._param = new._logits if new._logits is not None else new._probs
         new._num_events = self._num_events
-        super(Categorical, new).__init__(batch_shape, validate_args=False)
+        distribution.Distribution.__init__(
+            new, batch_shape, validate_args=False
+        )
         new._validate_args_enabled = self._validate_args_enabled
         return new
 

python/paddle/compat/nn/__init__.py

@@ -23,6 +23,7 @@
 from paddle.nn.modules.utils import _single
 from paddle.utils.decorator_utils import ForbidKeywordsDecorator
 
+from ..utils import _CompatClassMeta
 from . import functional
 from .transformer import MultiheadAttention
 
@@ -58,7 +59,7 @@
 ]
 
 
-class BatchNorm1D(nn.BatchNorm1D):
+class BatchNorm1D(nn.BatchNorm1D, metaclass=_CompatClassMeta):
     def __init__(
         self,
         num_features: int,
@@ -85,7 +86,7 @@ def __init__(
         self.momentum = momentum
 
 
-class BatchNorm2D(nn.BatchNorm2D):
+class BatchNorm2D(nn.BatchNorm2D, metaclass=_CompatClassMeta):
     def __init__(
         self,
         num_features: int,
@@ -112,7 +113,7 @@ def __init__(
         self.momentum = momentum
 
 
-class BatchNorm3D(nn.BatchNorm3D):
+class BatchNorm3D(nn.BatchNorm3D, metaclass=_CompatClassMeta):
     def __init__(
         self,
         num_features: int,
@@ -144,7 +145,7 @@ def __init__(
 BatchNorm3d = BatchNorm3D
 
 
-class AvgPool1D(nn.Layer):
+class AvgPool1D(nn.Layer, metaclass=_CompatClassMeta):
     r"""
     This operation applies a 1D average pooling over an input signal composed
     of several input planes, based on the input, output_size, return_mask parameters.
@@ -227,7 +228,7 @@ def __init__(
         ceil_mode: bool = False,
         count_include_pad: bool = True,
     ) -> None:
-        super().__init__()
+        nn.Layer.__init__(self)
         self.kernel_size = _single(kernel_size)
         self.stride = _single(stride if stride is not None else kernel_size)
         self.padding = _single(padding)
@@ -248,7 +249,7 @@ def extra_repr(self) -> str:
         return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
 
 
-class AvgPool2D(nn.Layer):
+class AvgPool2D(nn.Layer, metaclass=_CompatClassMeta):
     r"""
     This operation applies 2D average pooling over input features based on the input,
     and kernel_size, stride, padding parameters. Input(X) and Output(Out) are
@@ -345,7 +346,7 @@ def __init__(
         count_include_pad: bool = True,
         divisor_override: int | None = None,
     ):
-        super().__init__()
+        nn.Layer.__init__(self)
         self.kernel_size = kernel_size
         self.stride = stride if (stride is not None) else kernel_size
         self.padding = padding
@@ -368,7 +369,7 @@ def extra_repr(self) -> str:
         return f"kernel_size={self.kernel_size}, stride={self.stride}, padding={self.padding}"
 
 
-class AvgPool3D(nn.Layer):
+class AvgPool3D(nn.Layer, metaclass=_CompatClassMeta):
     """
 
     This operation applies 3D max pooling over input features based on the input,
@@ -452,7 +453,7 @@ def __init__(
         count_include_pad: bool = True,
         divisor_override: int | None = None,
     ) -> None:
-        super().__init__()
+        nn.Layer.__init__(self)
         self.kernel_size = kernel_size
         self.stride = stride if (stride is not None) else kernel_size
         self.padding = padding
@@ -481,7 +482,7 @@ def __setstate__(self, state):
         self.__dict__.setdefault("count_include_pad", True)
 
 
-class Unfold(nn.Unfold):
+class Unfold(nn.Unfold, metaclass=_CompatClassMeta):
     """
     A compatible version of paddle.nn.Unfold:
 
@@ -547,7 +548,7 @@ def to_list_if_necessary(x):
         )
 
 
-class Linear(nn.Layer):
+class Linear(nn.Layer, metaclass=_CompatClassMeta):
     r"""
 
     Python compatible fully-connected linear transformation layer. For each input :math:`X` ,
@@ -639,7 +640,7 @@ def __init__(
         device: PlaceLike | None = None,
         dtype: DTypeLike | None = None,
     ) -> None:
-        super().__init__()
+        nn.Layer.__init__(self)
         self._dtype = (
             self._helper.get_default_dtype() if dtype is None else dtype
         )
@@ -687,7 +688,7 @@ def reset_parameters(self) -> None:
             nn.init.uniform_(self.bias, -bound, bound)
 
 
-class Softmax(nn.Layer):
+class Softmax(nn.Layer, metaclass=_CompatClassMeta):
     r"""
     Softmax Activation.
 
@@ -814,7 +815,7 @@ class Softmax(nn.Layer):
         correct_name="paddle.nn.Softmax",
     )
     def __init__(self, dim: int | None = None) -> None:
-        super().__init__()
+        nn.Layer.__init__(self)
         self._dim = dim
         self._dtype = None
 
@@ -825,7 +826,7 @@ def extra_repr(self) -> str:
         return f"dim={self._dim}"
 
 
-class SmoothL1Loss(nn.Layer):
+class SmoothL1Loss(nn.Layer, metaclass=_CompatClassMeta):
     r"""
 
     PyTorch compatible version of :ref:`api_paddle_nn_SmoothL1Loss`, aligned with
@@ -890,7 +891,7 @@ def __init__(
         reduction: str = 'mean',
         beta: float = 1.0,
     ) -> None:
-        super().__init__()
+        nn.Layer.__init__(self)
         if size_average is not None or reduce is not None:
             reduction = (
                 'none'