chore: import upstream snapshot with attribution

python/paddle/jit/marker.py

@@ -0,0 +1,188 @@
+# 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 inspect
+from typing import TYPE_CHECKING, Protocol, TypeVar, overload
+
+from typing_extensions import (
+    ParamSpec,
+)
+
+import paddle
+from paddle.jit.dy2static.utils import DYNAMIC_DIMS_ATTR_NAME
+
+from .dy2static.utils import (
+    TransformOptions,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+
+_RetT = TypeVar("_RetT")
+_InputT = ParamSpec("_InputT")
+
+
+class _NotToStaticDecorator(Protocol):
+    @overload
+    def __call__(
+        self, func: Callable[_InputT, _RetT]
+    ) -> Callable[_InputT, _RetT]: ...
+
+    @overload
+    def __call__(self, func: None = ...) -> _NotToStaticDecorator: ...
+
+
+@overload
+def not_to_static(
+    func: Callable[_InputT, _RetT],
+) -> Callable[_InputT, _RetT]: ...
+
+
+@overload
+def not_to_static(func: None = ...) -> _NotToStaticDecorator: ...
+
+
+# Legacy decorator only for AST
+def not_to_static(func=None):
+    """
+    A Decorator to suppresses the convention of a function.
+
+    Args:
+        func(callable): The function to decorate.
+
+    Returns:
+        callable: A function which won't be converted in Dynamic-to-Static.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
+            >>> import paddle
+
+            >>> @paddle.jit.not_to_static
+            ... def func_not_to_static(x):
+            ...     res = x - 1
+            ...     return res
+
+            >>> @paddle.jit.to_static
+            ... def func(x):
+            ...     if paddle.mean(x) < 0:
+            ...         out = func_not_to_static(x)
+            ...     else:
+            ...         out = x + 1
+            ...     return out
+            >>> x = paddle.ones([1, 2], dtype='float32')
+            >>> out = func(x)
+            >>> print(out)
+            Tensor(shape=[1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[2., 2.]])
+    """
+    return unified(func, for_sot=False, for_ast=True)
+
+
+def unified(
+    fn: Callable[_InputT, _RetT] | type[paddle.nn.Layer] | None = None,
+    *,
+    for_sot: bool = True,
+    for_ast: bool = True,
+) -> Callable[_InputT, _RetT]:
+    """
+    Mark a function already unified in dygraph and static mode. So
+    that it won't be transformed again in SOT or AST mode.
+
+    Args:
+        fn(callable): The function to decorate.
+        for_sot(bool): Whether to mark the function as unified in SOT mode.
+        for_ast(bool): Whether to mark the function as unified in AST mode.
+    """
+
+    def _mark_as_unified(fn, *, for_sot: bool, for_ast: bool):
+        mode = TransformOptions.ToStaticMode.Nil()
+        if for_sot:
+            mode |= TransformOptions.ToStaticMode.SOT
+        if for_ast:
+            mode |= TransformOptions.ToStaticMode.AST
+        options = TransformOptions().with_skip_transform_mode(mode)
+        options.attach(fn)
+        return fn
+
+    if fn is None:
+        return lambda fn: _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
+    return _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
+
+
+def capture_control_flow(
+    fn: Callable[_InputT, _RetT] | None = None,
+) -> Callable[_InputT, _RetT]:
+    def _mark_as_need_capture_control_flow(fn):
+        options = TransformOptions().with_need_capture_control_flow(True)
+        options.attach(fn)
+        return fn
+
+    if fn is None:
+        return _mark_as_need_capture_control_flow
+    return _mark_as_need_capture_control_flow(fn)
+
+
+def force_dynamic(
+    fn: Callable[_InputT, _RetT] | type[paddle.nn.Layer] | None = None,
+) -> Callable[_InputT, _RetT]:
+    """
+    Mark a function or paddle.nn.Layer to be executed in dynamic mode, it will
+    break the graph and prevent it from being converted to static mode.
+    """
+    from paddle.jit import sot
+
+    if inspect.isclass(fn) and issubclass(fn, paddle.nn.Layer):
+        sot.utils.paddle_api_config.add_break_graph_layer_class(fn)
+        return fn
+    if inspect.isfunction(fn):
+        sot.utils.paddle_api_config.add_break_graph_function(fn)
+        return fn
+
+    raise TypeError(
+        f"Expected a callable or paddle.nn.Layer, but got {type(fn).__name__}."
+    )
+
+
+def dynamic_dims(
+    tensor: paddle.Tensor,
+    dims: int | Sequence[int],
+) -> None:
+    """
+    Mark a tensor as having dynamic dimensions.
+    This is used to indicate that the tensor's shape may change dynamically
+    during execution, which is particularly useful in dynamic-to-static
+    conversion scenarios.
+    Args:
+        tensor (paddle.Tensor): The tensor to mark as dynamic.
+        dims (int | Sequence[int]): The dimensions to mark as dynamic.
+    """
+    if not isinstance(tensor, paddle.Tensor):
+        raise TypeError(
+            f"Expected a paddle.Tensor, but got {type(tensor).__name__}."
+        )
+
+    if isinstance(dims, int):
+        dims = (dims,)
+
+    if not isinstance(dims, (list, tuple)):
+        raise TypeError("Dimensions must be a list or tuple.")
+    if not all(isinstance(dim, int) for dim in dims):
+        raise TypeError("All dimensions must be integers.")
+
+    setattr(tensor, DYNAMIC_DIMS_ATTR_NAME, tuple(dims))