# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
# Copyright (c) 2021 NVIDIA Corporation. 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
import os
import pickle
import sys
import threading
import types
import warnings
from collections import OrderedDict
from collections.abc import Sequence
from contextlib import contextmanager
from types import ModuleType
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Protocol,
    TypedDict,
    TypeVar,
    overload,
)

from typing_extensions import (
    Literal,
    NotRequired,
    ParamSpec,
    TypeAlias,
    Unpack,
)

import paddle
from paddle.base import core, dygraph
from paddle.base.compiler import (
    BuildStrategy,
)
from paddle.base.dygraph.base import (
    switch_to_static_graph,
)
from paddle.base.executor import Executor, scope_guard
from paddle.base.framework import (
    EagerParamBase,
    Parameter,
    Variable,
    _current_expected_place,
    dygraph_only,
)
from paddle.base.wrapped_decorator import wrap_decorator
from paddle.framework import use_pir_api
from paddle.nn import Layer
from paddle.static.io import save_inference_model
from paddle.utils.environments import (
    EnvironmentVariableGuard,
)

from .dy2static import logging_utils
from .dy2static.convert_call_func import add_ignore_module
from .dy2static.program_translator import (
    ASTStaticFunction,
    ProgramTranslator,
    StaticFunction,
    SymbolicStaticFunction,
    unwrap_decorators,
)
from .dy2static.utils import (
    ENV_ENABLE_SOT,
    Backend,
    infer_use_cinn_backend,
)
from .pir_translated_layer import PIR_INFER_MODEL_SUFFIX, PirTranslatedLayer
from .translated_layer import (
    INFER_MODEL_SUFFIX,
    INFER_PARAMS_INFO_SUFFIX,
    INFER_PARAMS_SUFFIX,
    INFER_PROPERTY_SUFFIX,
    TranslatedLayer,
)

if TYPE_CHECKING:
    from paddle import Tensor
    from paddle._typing import NestedStructure
    from paddle.static import InputSpec

    class _SaveOptions(TypedDict):
        output_spec: NotRequired[Sequence[Tensor | int]]
        with_hook: NotRequired[bool]
        combine_params: NotRequired[bool]
        clip_extra: NotRequired[bool]
        skip_forward: NotRequired[bool]
        input_names_after_prune: NotRequired[list[str]]
        skip_prune_program: NotRequired[bool]
        separate_parameters: NotRequired[bool]

    class _LoadOptions(TypedDict):
        model_filename: NotRequired[str]
        params_filename: NotRequired[str]


_LayerT = TypeVar("_LayerT", bound=Layer)
_RetT = TypeVar("_RetT")
_InputT = ParamSpec("_InputT")
Backends: TypeAlias = Literal["CINN"]


@contextmanager
def sot_mode_guard(value: bool):
    with EnvironmentVariableGuard(ENV_ENABLE_SOT, value):
        yield


def copy_decorator_attrs(original_func, decorated_obj):
    """
    Copies some necessary attributes from original function into decorated function.

    Args:
        original_func(callable): the original decorated function.
        decorated_obj(StaticFunction): the target decorated StaticFunction object.
    """
    decorator_name = "to_static"

    decorated_obj.__name__ = original_func.__name__
    decorated_obj._decorator_name = decorator_name
    if not inspect.ismethod(original_func):
        decorated_obj.__wrapped__ = original_func
    decorated_obj.__doc__ = original_func.__doc__
    if hasattr(original_func, "__module__"):
        decorated_obj.__module__ = original_func.__module__

    return decorated_obj


def ignore_module(modules: list[ModuleType]) -> None:
    """
    Adds modules that ignore transcription.
    Builtin modules that have been ignored are collections, pdb, copy, inspect, re, numpy, logging, six

    Args:
        modules (list[ModuleType]): Ignored modules that you want to add

    Examples:
        .. code-block:: pycon

            >>> import scipy
            >>> import networkx

            >>> import paddle
            >>> from paddle.jit import ignore_module
            >>> modules = [
            ...     scipy,
            ...     networkx,
            ... ]
            >>> ignore_module(modules)

    """
    add_ignore_module(modules)


class _ToStaticOptions(TypedDict):
    property: NotRequired[bool]
    full_graph: NotRequired[bool]


class _ToStaticDecorator(Protocol):
    @overload
    def __call__(self, function: _LayerT) -> _LayerT: ...

    @overload
    def __call__(
        self, function: Callable[_InputT, _RetT]
    ) -> StaticFunction[_InputT, _RetT]: ...


@overload
def to_static(
    function: _LayerT,
    input_spec: NestedStructure[InputSpec] | None = ...,
    build_strategy: BuildStrategy | None = ...,
    backend: Backends | None = ...,
    **kwargs: Unpack[_ToStaticOptions],
) -> _LayerT: ...


@overload
def to_static(
    function: Callable[_InputT, _RetT],
    input_spec: NestedStructure[InputSpec] | None = ...,
    build_strategy: BuildStrategy | None = ...,
    backend: Backends | None = ...,
    **kwargs: Unpack[_ToStaticOptions],
) -> StaticFunction[_InputT, _RetT]: ...


@overload
def to_static(
    function: None = ...,
    input_spec: NestedStructure[InputSpec] | None = ...,
    build_strategy: BuildStrategy | None = ...,
    backend: Backends | None = ...,
    **kwargs: Unpack[_ToStaticOptions],
) -> _ToStaticDecorator: ...


def to_static(
    function=None,
    input_spec=None,
    build_strategy=None,
    backend="CINN",
    **kwargs,
):
    """
    Converts dynamic graph APIs into static graph function APIs. Decorator
    @to_static handles the Program and Executor of static graph mode and returns
    the result as dynamic graph Tensor(s). Users could use the returned dynamic
    graph Tensor(s) to do dynamic graph training, inference, or other operations.
    If the decorated function calls other dynamic graph function, the called one
    will be converted into static graph function as well.

    Args:
        function (callable): Callable dynamic graph function. If it used as a
            decorator, the decorated function will be parsed as this parameter.
        input_spec (list[InputSpec]|tuple[InputSpec]): list/tuple of InputSpec to
            specific the shape/dtype/name information of each input Tensor.
        build_strategy (BuildStrategy|None): This argument is used to compile the
            converted program with the specified options, such as operators' fusion
            in the computational graph and memory optimization during the execution
            of the computational graph. For more information about :attr:`build_strategy`,
            please refer to :ref:`paddle.static.BuildStrategy <cn_api_paddle_static_BuildStrategy>`.
            The default is ``None``.
        backend(str, Optional): Specifies compilation backend, which can be ``"CINN"`` or
            ``None``. When backend is ``"CINN"``, CINN compiler will be used to speed up
            training and inference. default value is ``"CINN"``.
        kwargs: Support keys including :attr:`property` and :attr:`full_graph`.

          - property (bool): If True, the function will be treated as a property
            function. The default is False.
          - full_graph (bool): If True, the function will be converted into a
            full static graph. The default is False.

    Returns:
        Tensor(s): containing the numerical result.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
            >>> import paddle
            >>> from paddle.jit import to_static

            >>> @to_static
            >>> def func(x):
            ...     if paddle.mean(x) < 0:
            ...         x_v = x - 1
            ...     else:
            ...         x_v = x + 1
            ...     return x_v
            >>> x = paddle.ones([1, 2], dtype='float32')
            >>> x_v = func(x)
            >>> print(x_v)
            Tensor(shape=[1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
            [[2., 2.]])

    """
    property = kwargs.get("property", False)
    full_graph = kwargs.get("full_graph", None)
    build_strategy = build_strategy or BuildStrategy()
    if not isinstance(build_strategy, BuildStrategy):
        raise TypeError(
            f"Required type(build_strategy) shall be `paddle.static.BuildStrategy`, but received {type(build_strategy).__name__}"
        )
    backend = Backend.from_arg(backend)
    if infer_use_cinn_backend(backend, build_strategy):
        backend = Backend.CINN
    elif backend.is_pcc():
        pass
    else:
        backend = Backend.PHI

    def decorated(python_func):
        """
        Decorates a python function into a ASTStaticFunction or SymbolicStaticFunction object.
        """

        nonlocal full_graph
        if full_graph is None:
            flag = ENV_ENABLE_SOT.get()
            full_graph = not flag

        if sys.version_info >= (3, 15) and not full_graph:
            warnings.warn(
                "full_graph=False is not supported in Python 3.15+. Set full_graph=True automatically"
            )
            full_graph = True

        StaticClass = (
            ASTStaticFunction if full_graph else SymbolicStaticFunction
        )

        # Step 1. unwrap the function if it is already decorated.
        _, python_func = unwrap_decorators(python_func)

        # Step 2. copy some attributes from original python function.
        static_layer = copy_decorator_attrs(
            original_func=python_func,
            decorated_obj=StaticClass(
                function=python_func,
                input_spec=input_spec,
                build_strategy=build_strategy,
                property=property,
                backend=backend,
            ),
        )

        return static_layer

    # for usage: `to_static(foo, ...)`
    if function is not None:
        if isinstance(function, Layer):
            if isinstance(function.forward, StaticFunction):
                class_name = function.__class__.__name__
                logging_utils.warn(
                    f"`{class_name}.forward` has already been decorated somewhere. It will be redecorated to replace previous one."
                )
            function._original_funcs["forward"] = function.forward
            function.forward = decorated(function.forward)
            return function
        else:
            return decorated(function)

    # for usage: `@to_static`
    return decorated


class _SaveLoadConfig:
    def __init__(self):
        self._output_spec = None
        self._model_filename = None
        self._params_filename = None
        self._separate_params = False
        # used for `paddle.load`
        self._keep_name_table = False

        # NOTE: Users rarely use following configs, so these configs are not open to users,
        # reducing user learning costs, but we retain the configuration capabilities

        # If True, programs are modified to only support direct inference deployment.
        # Otherwise,more information will be stored for flexible optimization and re-training.
        # Currently, only True is supported
        self._export_for_deployment = True

        # If True, It will save inference program only, and do not save params of Program
        self._program_only = False
        self.with_hook = False

        # if True, multi `StaticFunction` will share params in one file.
        self.combine_params = False

        # when need to save a prune model, use input_names_after_prune to specify the inputs left after pruning
        self.input_names_after_prune = None

        # in the scene of llm-inference, pruning program can cause unexpectable result, an option to skip prune is necessary
        self.skip_prune_program = False

        # if True, the params will be saved separately in multiple files.
        self.separate_parameters = False

    @property
    def output_spec(self):
        return self._output_spec

    @output_spec.setter
    def output_spec(self, spec):
        if spec is None:
            return
        if not isinstance(spec, list):
            raise TypeError(
                f"The config `output_spec` should be 'list', but received input type is {type(input)}."
            )
            for var in spec:
                if not isinstance(var, core.eager.Tensor):
                    raise TypeError(
                        f"The element in config `output_spec` list should be 'Variable', but received element's type is {type(var)}."
                    )
        self._output_spec = spec

    @property
    def model_filename(self):
        return self._model_filename

    @model_filename.setter
    def model_filename(self, filename):
        if filename is None:
            return
        if not isinstance(filename, str):
            raise TypeError(
                f"The config `model_filename` should be str, but received input's type is {type(filename)}."
            )
        if len(filename) == 0:
            raise ValueError("The config `model_filename` is empty string.")
        self._model_filename = filename

    @property
    def params_filename(self):
        return self._params_filename

    @params_filename.setter
    def params_filename(self, filename):
        if filename is None:
            return
        if not isinstance(filename, str):
            raise TypeError(
                f"The config `params_filename` should be str, but received input's type is {type(filename)}."
            )
        if len(filename) == 0:
            raise ValueError("The config `params_filename` is empty string.")
        self._params_filename = filename

    @property
    def keep_name_table(self):
        return self._keep_name_table

    @keep_name_table.setter
    def keep_name_table(self, value):
        if value is None:
            return
        if not isinstance(value, bool):
            raise TypeError(
                f"The config `keep_name_table` should be bool value, but received input's type is {type(value)}."
            )
        self._keep_name_table = value


def _parse_save_configs(configs: _SaveOptions) -> _SaveLoadConfig:
    supported_configs = [
        "output_spec",
        "with_hook",
        "combine_params",
        "clip_extra",
        "skip_forward",
        "input_names_after_prune",
        "skip_prune_program",
        "separate_parameters",
    ]

    # input check
    for key in configs:
        if key not in supported_configs:
            raise ValueError(
                f"The additional config ({key}) of `paddle.jit.save` is not supported."
            )

    # construct inner config
    inner_config = _SaveLoadConfig()
    inner_config.output_spec = configs.get("output_spec", None)
    inner_config.with_hook = configs.get("with_hook", False)
    inner_config.combine_params = configs.get("combine_params", False)
    inner_config.clip_extra = configs.get("clip_extra", True)
    inner_config.skip_forward = configs.get("skip_forward", False)
    inner_config.input_names_after_prune = configs.get(
        "input_names_after_prune", None
    )
    inner_config.skip_prune_program = configs.get("skip_prune_program", False)
    inner_config.separate_parameters = configs.get("separate_parameters", False)

    return inner_config


def _parse_load_config(configs: _LoadOptions) -> _SaveLoadConfig:
    supported_configs = ['model_filename', 'params_filename']

    # input check
    for key in configs:
        if key not in supported_configs:
            raise ValueError(
                f"The additional config ({key}) of `paddle.jit.load` is not supported."
            )

    # construct inner config
    inner_config = _SaveLoadConfig()
    inner_config.model_filename = configs.get('model_filename', None)
    inner_config.params_filename = configs.get('params_filename', None)

    return inner_config


def _get_input_var_and_names(inputs, input_spec, input_names_after_prune):
    name_none_error = (
        "The %s's name is None. "
        "When using jit.save, please set InputSpec's name in "
        "to_static(input_spec=[]) and jit.save(input_spec=[]) "
        "and make sure they are consistent."
    )
    name_no_exists_error = (
        "The tensor `%s` does not exists. "
        "Please make sure the name of InputSpec or example Tensor "
        "in input_spec is the same as the name of InputSpec in "
        "`to_static` decorated on the Layer.forward method."
    )
    if input_names_after_prune is not None:
        input_spec = [
            x
            for x in input_spec
            if isinstance(x, paddle.static.InputSpec)
            and x.name in input_names_after_prune
        ]

    input_vars = [
        var
        for var in paddle.utils.flatten(inputs)
        if isinstance(var, (Variable, paddle.pir.Value))
    ]
    input_var_names = [
        var.name
        for var in paddle.utils.flatten(inputs)
        if isinstance(var, (Variable, paddle.pir.Value))
    ]

    if input_spec is None:
        # no prune
        return input_vars, input_var_names
    else:
        # filter out non-tensor type spec infos.
        input_spec = [
            spec
            for spec in input_spec
            if isinstance(spec, paddle.static.InputSpec)
        ]
    result_var_list = []
    result_name_list = []
    if len(input_spec) == len(input_var_names):
        # no prune
        result_var_list = input_vars
        result_name_list = input_var_names
        # if input spec name not in input_var_names, only raise warning
        for spec in input_spec:
            if spec.name is None:
                warnings.warn(name_none_error % spec)
            elif spec.name not in input_var_names:
                warnings.warn(name_no_exists_error % spec.name)
            else:
                pass
    else:
        # prune
        for spec in input_spec:
            if spec.name is None:
                # name is None, the input_spec only can be InputSpec
                raise ValueError(name_none_error % spec)
            elif spec.name not in input_var_names:
                # the input_spec can be `InputSpec` or `Tensor`
                raise ValueError(name_no_exists_error % spec.name)
            else:
                result_var_list.append(spec)
                result_name_list.append(spec.name)

    return result_var_list, result_name_list


def _contains_dict(output):
    if isinstance(output, dict):
        return True
    if isinstance(output, Sequence) and not isinstance(output, str):
        return any(_contains_dict(i) for i in output)
    return False


def _get_output_vars(outputs, output_spec, with_hook=False):
    name_no_exists_error = (
        "The tensor `%s` does not exists. "
        "Please make sure the name of example Tensor "
        "in configs.output_spec is the output tensor of "
        "Layer.forward method."
    )
    output_spec_is_not_value_error = (
        "tensor `%s` is not support in pir mode, "
        "because pir value has no name sometimes, especially as output,"
        "so we can't check tensor's name with output var name, please"
        "change as pir.value(to_static layer's output)"
        "or int(the position of to_static layer's output)"
    )
    if output_spec and with_hook:
        raise RuntimeError(
            "Currently not support specify output_spec while founding pre/post hooks in your outermost layer."
        )
    result_list = []
    if _contains_dict(outputs):
        warnings.warn(
            "Found 'dict' in given outputs, the values will be returned in a sequence sorted in lexicographical order by their keys."
        )
    if use_pir_api():
        from paddle.autograd.backward_utils import ValueSet

        for var in paddle.utils.flatten(outputs):
            if isinstance(var, paddle.pir.Value) and var not in ValueSet(
                result_list
            ):
                result_list.append(var)

        if output_spec is not None:
            output_size = len(result_list)
            if len(output_spec) == output_size:
                for var in output_spec:
                    if not isinstance(var, (paddle.pir.Value, int)):
                        warnings.warn(output_spec_is_not_value_error % var.name)
                    else:
                        if var not in ValueSet(result_list):
                            warnings.warn(name_no_exists_error % var.name)
            else:
                result_set = ValueSet(result_list)
                part_result_list = []
                for var in output_spec:
                    if isinstance(var, paddle.pir.Value):
                        if var not in result_set:
                            raise ValueError(name_no_exists_error % var.name)
                        else:
                            part_result_list.append(var)
                    elif isinstance(var, int):
                        if var >= output_size:
                            raise ValueError(
                                "position %d should smaller than output's size % d",
                                var,
                                output_size,
                            )
                        else:
                            part_result_list.append(result_list[var])

                    else:
                        raise ValueError(
                            output_spec_is_not_value_error % var.name
                        )

                return part_result_list
    else:
        output_vars_dict = OrderedDict()
        for var in paddle.utils.flatten(outputs):
            if isinstance(var, (Variable)):
                output_vars_dict[var.name] = var
        if output_spec is None:
            result_list = list(output_vars_dict.values())
        elif output_spec is not None and len(output_spec) == len(
            output_vars_dict
        ):
            result_list = list(output_vars_dict.values())
            for var in output_spec:
                if var.name not in output_vars_dict:
                    warnings.warn(name_no_exists_error % var.name)
        else:
            for var in output_spec:
                if var.name not in output_vars_dict:
                    raise ValueError(name_no_exists_error % var.name)
                else:
                    result_list.append(output_vars_dict[var.name])

    return result_list


# NOTE(chenweihang): [ Handling of use cases of API paddle.jit.load ]
# `paddle.jit.load` may be used to load saved results of:
# 1. Expected cases:
#   - paddle.jit.save
#   - paddle.static.save_inference_model
# 2. Error cases:
#   - paddle.save: no .pdmodel for prefix
#   - paddle.static.save: no .pdiparams but .pdparams exists
#   - paddle.base.io.save_params/save_persistables: no __model__
# TODO(chenweihang): polish error message in above error cases
def _build_load_path_and_config(path, config):
    # NOTE(chenweihang): If both [prefix save format] and [directory save format] exist,
    # raise error, avoid confusing behavior
    if use_pir_api():
        model_suffix = PIR_INFER_MODEL_SUFFIX
    else:
        model_suffix = INFER_MODEL_SUFFIX
    prefix_format_path = path + model_suffix
    prefix_format_exist = os.path.exists(prefix_format_path)
    directory_format_exist = os.path.isdir(path)
    if prefix_format_exist and directory_format_exist:
        raise ValueError(
            f"The {path}.pdmodel(json) and {path} directory exist at the same time, "
            "don't know which one to load, please make sure that the specified target "
            "of ``path`` is unique."
        )
    elif not prefix_format_exist and not directory_format_exist:
        raise ValueError(
            f"The ``path`` ({path}) to load model not exists. "
            "Please make sure that *.pdmodel(json) exists or "
            "don't using ``skip_forward=True`` to jit.save."
        )
    else:
        if prefix_format_exist:
            file_prefix = os.path.basename(path)
            model_path = os.path.dirname(path)
            if config.model_filename is not None:
                warnings.warn(
                    "When loading the result saved with the "
                    "specified file prefix, the ``model_filename`` config does "
                    "not take effect."
                )
            config.model_filename = file_prefix + model_suffix
            if config.params_filename is not None:
                warnings.warn(
                    "When loading the result saved with the "
                    "specified file prefix, the ``params_filename`` config does "
                    "not take effect."
                )
            config.params_filename = file_prefix + INFER_PARAMS_SUFFIX
        else:
            # Compatible with the old save_inference_model format
            model_path = path

    return model_path, config


_save_pre_hooks_lock = threading.Lock()
_save_pre_hooks = []


class HookRemoveHelper:
    """A HookRemoveHelper that can be used to remove hook."""

    def __init__(self, hook):
        self._hook = hook

    def remove(self):
        _remove_save_pre_hook(self._hook)


def _register_save_pre_hook(hook):
    """
    Register a save pre-hook for `paddle.jit.save`.
    This hook will be executed before `save` function has been invoked.

    hook(layer, input_spec, configs) -> None
    - layer (Layer|function): This argument is corresponding to `layer` in `paddle.jit.save`.
    - input_spec (list or tuple[InputSpec|Tensor|Python built-in variable]): This argument is corresponding to `input_spec` in `paddle.jit.save`.
    - configs (dict): This argument is corresponding to `configs` in `paddle.jit.save`.

    Args:
        hook(function): a function registered as a save pre-hook

    Returns:
        HookRemoveHelper: a HookRemoveHelper object that can be used to remove the added hook by calling `hook_remove_helper.remove()`.

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('`paddle.jit.api.to_static` can not run in xdoctest')
            >>> import numpy as np
            >>> import paddle

            >>> IMAGE_SIZE = 256
            >>> CLASS_NUM = 10

            >>> class LinearNet(paddle.nn.Layer):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self._linear = paddle.nn.Linear(IMAGE_SIZE, CLASS_NUM)
            ...
            ...     def forward(self, x):
            ...         return self._linear(x)
            >>> saving_count = 0
            >>> def save_pre_hook(layer, input_spec, configs):
            ...     global saving_count
            ...     saving_count += 1
            >>> remove_handler = paddle.jit.api._register_save_pre_hook(save_pre_hook)

            >>> layer = LinearNet()
            >>> paddle.jit.save(layer, "/tmp", [paddle.static.InputSpec(shape=[-1, IMAGE_SIZE])])
            >>> print(saving_count)
            1

            >>> remove_handler.remove()
            >>> paddle.jit.save(layer, "/tmp", [paddle.static.InputSpec(shape=[-1, IMAGE_SIZE])])
            >>> print(saving_count)
            1
    """
    global _save_pre_hooks_lock
    global _save_pre_hooks
    _save_pre_hooks_lock.acquire()
    if hook not in _save_pre_hooks:
        _save_pre_hooks.append(hook)
    _save_pre_hooks_lock.release()
    return HookRemoveHelper(hook)


def _clear_save_pre_hooks():
    global _save_pre_hooks_lock
    global _save_pre_hooks
    _save_pre_hooks_lock.acquire()
    _save_pre_hooks.clear()
    _save_pre_hooks_lock.release()


def _remove_save_pre_hook(hook):
    global _save_pre_hooks_lock
    global _save_pre_hooks
    _save_pre_hooks_lock.acquire()
    if hook in _save_pre_hooks:
        _save_pre_hooks.remove(hook)
    _save_pre_hooks_lock.release()


class _SaveFunction(Protocol):
    def __call__(
        self,
        layer: Layer | Callable[..., Any],
        path: str,
        input_spec: Sequence[InputSpec | Tensor | object] | None = ...,
        **configs: Unpack[_SaveOptions],
    ) -> None: ...


@wrap_decorator
def _run_save_pre_hooks(func: _SaveFunction) -> _SaveFunction:
    def wrapper(
        layer: Layer | Callable[..., Any],
        path: str,
        input_spec: Sequence[InputSpec | Tensor | object] | None = None,
        **configs: Unpack[_SaveOptions],
    ) -> None:
        global _save_pre_hooks
        for hook in _save_pre_hooks:
            hook(layer, input_spec, configs)
        func(layer, path, input_spec, **configs)

    return wrapper


def _save_property(filename: str, property_vals: list[tuple[Any, str]]):
    """class property serialization.

    Args:
        filename (str): *.meta
        property_vals (list[tuple[Any, str]]): class property.
    """

    def set_property(meta, key, val):
        if isinstance(val, float):
            meta.set_float(key, val)
        elif isinstance(val, int):
            meta.set_int(key, val)
        elif isinstance(val, str):
            meta.set_string(key, val)
        elif isinstance(val, (tuple, list)):
            if isinstance(val[0], float):
                meta.set_floats(key, val)
            elif isinstance(val[0], int):
                meta.set_ints(key, val)
            elif isinstance(val[0], str):
                meta.set_strings(key, val)
        else:
            raise ValueError(f"Note support val type: {type(val)}")

    with open(filename, 'wb') as f:
        meta = paddle.framework.core.Property()
        for item in property_vals:
            val, key = item[0], item[1]
            set_property(meta, key, val)
        f.write(meta.serialize_to_string())


def _get_function_names_from_layer(layer: Layer) -> list[str]:
    cls = layer.__class__
    return [
        member_name
        for member_name, member in inspect.getmembers(cls)
        if (
            inspect.isfunction(member)
            or inspect.ismethod(member)
            or inspect.ismethoddescriptor(member)
        )
    ]


@_run_save_pre_hooks
@switch_to_static_graph
def save(
    layer: Layer | Callable[..., Any],
    path: str,
    input_spec: Sequence[InputSpec | Tensor | object] | None = None,
    **configs: Unpack[_SaveOptions],
) -> None:
    """
    Saves input Layer or function as ``paddle.jit.TranslatedLayer``
    format model, which can be used for inference or fine-tuning after loading.

    It will save the translated program and all related persistable
    variables of input Layer to given ``path`` .

    ``path`` is the prefix of saved objects, and the saved translated program file
    suffix is ``.pdmodel`` , the saved persistable variables file suffix is ``.pdiparams`` ,
    and here also saved some additional variable description information to a file,
    its suffix is ``.pdiparams.info``, these additional information is used in fine-tuning.

    The saved model can be loaded by follow APIs:
      - ``paddle.jit.load``
      - ``paddle.static.load_inference_model``
      - Other C++ inference APIs

    .. note::
        When using ``paddle.jit.save`` to save a function, parameters will not be saved. If you have to
        save the parameter, please pass the Layer containing function and parameter to ``paddle.jit.save``.

    Args:
        layer (Layer|function): The Layer or function to be saved.
        path (str): The path prefix to save model. The format is ``dirname/file_prefix`` or ``file_prefix``.
        input_spec (list or tuple[InputSpec|Tensor|Python built-in variable], optional): Describes the input of the saved model's forward
            method, which can be described by InputSpec or example Tensor. Moreover, we support to specify non-tensor type argument,
            such as int, float, string, or list/dict of them.If None, all input variables of
            the original Layer's forward method would be the inputs of the saved model. Default None.
        **configs (dict, optional): Other save configuration options for compatibility. We do not
            recommend using these configurations, they may be removed in the future. If not necessary,
            DO NOT use them. Default None.
            The following options are currently supported:
            (1) output_spec (list[Tensor|Value|int]): Selects the output targets of the saved model,
            By default, all return variables of original Layer's forward method are kept as the
            output of the saved model. If the provided ``output_spec`` list is not all output variables,
            the saved model will be pruned according to the given ``output_spec`` list.
            in pir mode, Tensor is not supported, because value has no name in most cases,
            which can't be used to judge which tensor corresponds to which value; the value can't be found
            if the saved program is not the same as the program that includes output_spec, so we need to
            use the position of the output.

    Returns:
        None

    Examples:
        .. code-block:: pycon

            >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
            >>> # example 1: save layer
            >>> import numpy as np
            >>> import paddle
            >>> import paddle.nn as nn
            >>> import paddle.optimizer as opt

            >>> BATCH_SIZE = 16
            >>> BATCH_NUM = 4
            >>> EPOCH_NUM = 4

            >>> IMAGE_SIZE = 784
            >>> CLASS_NUM = 10

            >>> # define a random dataset
            >>> class RandomDataset(paddle.io.Dataset):  # type: ignore[type-arg]
            ...     def __init__(self, num_samples):
            ...         self.num_samples = num_samples
            ...
            ...     def __getitem__(self, idx):
            ...         image = np.random.random([IMAGE_SIZE]).astype('float32')
            ...         label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
            ...         return image, label
            ...
            ...     def __len__(self):
            ...         return self.num_samples

            >>> class LinearNet(nn.Layer):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
            ...
            ...     @paddle.jit.to_static
            ...     def forward(self, x):
            ...         return self._linear(x)

            >>> def train(layer, loader, loss_fn, opt):
            ...     for epoch_id in range(EPOCH_NUM):
            ...         for batch_id, (image, label) in enumerate(loader()):
            ...             out = layer(image)
            ...             loss = loss_fn(out, label)
            ...             loss.backward()
            ...             opt.step()
            ...             opt.clear_grad()
            ...             print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))

            >>> # 1. train & save model.

            >>> # create network
            >>> layer = LinearNet()
            >>> loss_fn = nn.CrossEntropyLoss()
            >>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())

            >>> # create data loader
            >>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
            >>> loader = paddle.io.DataLoader(
            ...     dataset,
            ...     batch_size=BATCH_SIZE,
            ...     shuffle=True,
            ...     drop_last=True,
            ...     num_workers=2,
            ... )

            >>> # train
            >>> train(layer, loader, loss_fn, adam)

            >>> # save
            >>> path = "example_model/linear"
            >>> paddle.jit.save(layer, path)

            >>> # example 2: save function
            >>> import paddle
            >>> from paddle.static import InputSpec


            >>> def save_function():
            ...     @paddle.jit.to_static
            ...     def fun(inputs):
            ...         return paddle.tanh(inputs)
            ...
            ...     path = 'test_jit_save_load_function_1/func'
            ...     inps = paddle.rand([3, 6])
            ...     origin = fun(inps)
            ...
            ...     paddle.jit.save(fun, path)
            ...     load_func = paddle.jit.load(path)
            ...
            ...     load_result = load_func(inps)
            ...     print((load_result - origin).abs().max() < 1e-10)

            >>> save_function()
    """
    # 1. input build & check
    prog_translator = ProgramTranslator()
    is_prim_infer = core._is_fwd_prim_enabled() and core._is_bwd_prim_enabled()
    if not prog_translator.enable_to_static:
        raise RuntimeError(
            "The paddle.jit.save doesn't work when setting 'paddle.jit.enable_to_static' to False."
        )

    if not (
        isinstance(layer, (Layer, StaticFunction)) or inspect.isfunction(layer)
    ):
        raise TypeError(
            f"The input of paddle.jit.save should be 'Layer' or 'Function', but received input type is {type(layer)}."
        )
    elif inspect.isfunction(layer) or isinstance(layer, StaticFunction):
        warnings.warn(
            'What you save is a function, and `jit.save` will generate the name of the model file according to `path` you specify. When loading these files with `jit.load`, you get a `TranslatedLayer` whose inference result is the same as the inference result of the function you saved.'
        )

    # NOTE(chenweihang): If the input layer be wrapped by DataParallel,
    # the args and kwargs of forward method will can't be parsed by
    # function_spec, so here we save DataParallel._layers instead
    # DataParallel it self
    # NOTE(chenweihang): using inner_layer, do not change input layer
    if isinstance(layer, paddle.DataParallel):
        inner_layer = layer._layers
    else:
        inner_layer = layer

    # path check
    file_prefix = os.path.basename(path)
    if file_prefix == "":
        raise ValueError(
            "The input path MUST be format of dirname/file_prefix "
            "[dirname\\file_prefix in Windows system], but received "
            "file_prefix is empty string."
        )

    dirname = os.path.dirname(path)
    if dirname and not os.path.exists(dirname):
        os.makedirs(dirname)

    # avoid change user given input_spec
    inner_input_spec = None
    if input_spec is not None:
        if isinstance(layer, Layer):
            for member_name in _get_function_names_from_layer(inner_layer):
                static_func = getattr(inner_layer, member_name, None)
                if (
                    isinstance(static_func, StaticFunction)
                    and 'forward' != member_name
                ):
                    raise ValueError(
                        f"If there are static functions other than 'forward' that need to be saved, the input 'input_spec' should be None, but received the type of 'input_spec' is {type(input_spec)}."
                    )

        if not isinstance(input_spec, (list, tuple)):
            raise TypeError(
                f"The input input_spec should be 'list', but received input_spec's type is {type(input_spec)}."
            )
        inner_input_spec = []
        for var in paddle.utils.flatten(input_spec):
            if isinstance(var, paddle.static.InputSpec):
                inner_input_spec.append(var)
            elif isinstance(
                var, (core.eager.Tensor, Variable, paddle.pir.Value)
            ):
                inner_input_spec.append(
                    paddle.static.InputSpec.from_tensor(var)
                )
            else:
                # NOTE(Aurelius84): Support non-Tensor type in `input_spec`.
                inner_input_spec.append(var)

    # parse configs
    configs = _parse_save_configs(configs)
    # whether outermost layer has pre/post hook, if does, we need also save
    # these operators in program.
    with_hook = configs.with_hook
    combine_params = configs.combine_params
    if combine_params:
        configs._program_only = True

    scope = core.Scope()
    extra_var_info = {}
    if isinstance(layer, Layer):
        functions = list(set(_get_function_names_from_layer(inner_layer)))
        functions = sorted(functions)
        if inner_layer._forward_pre_hooks or inner_layer._forward_post_hooks:
            with_hook = True
    else:
        # layer is function
        functions = [layer]

    combine_vars = {}
    combine_program = []
    property_vals = []  # (value, key)
    concrete_program = None
    for attr_func in functions:
        if isinstance(layer, Layer):
            static_func = get_ast_static_function(
                getattr(inner_layer, attr_func, None)
            )
            if isinstance(static_func, StaticFunction):
                if static_func.is_property:
                    # property method to be exported
                    immediate_val = static_func()
                    property_vals.append(
                        (
                            immediate_val,
                            layer.__class__.__name__ + '.' + attr_func,
                        )
                    )
                    continue

                concrete_program = (
                    static_func.concrete_program_specify_input_spec(
                        inner_input_spec,
                        with_hook=with_hook,
                        is_prim_infer=is_prim_infer,
                    )
                )

            elif 'forward' == attr_func:
                if configs.skip_forward:
                    # do not jit.save forward function
                    continue

                # transform in jit.save, if input_spec is incomplete, declarative will throw error
                # inner_input_spec is list[InputSpec], it should be packed with same structure
                # as original input_spec here.
                if inner_input_spec:
                    inner_input_spec = paddle.utils.pack_sequence_as(
                        input_spec, inner_input_spec
                    )
                static_forward = to_static(
                    inner_layer.forward,
                    input_spec=inner_input_spec,
                    full_graph=True,
                    backend=None,
                )

                concrete_program = (
                    static_forward.concrete_program_specify_input_spec(
                        with_hook=with_hook, is_prim_infer=is_prim_infer
                    )
                )
                # the input_spec has been used in declarative, which is equal to
                # @to_static with input_spec and jit.save without input_spec,
                # avoid needless warning
                inner_input_spec = None
            else:
                continue
        else:
            # When layer is a function
            if isinstance(attr_func, StaticFunction):
                static_func = get_ast_static_function(attr_func)

                if static_func.is_property:
                    # property method to be exported
                    immediate_val = static_func()
                    property_vals.append((immediate_val, static_func))
                    continue

                concrete_program = (
                    static_func.concrete_program_specify_input_spec(
                        inner_input_spec, is_prim_infer=is_prim_infer
                    )
                )
            else:
                static_func = get_ast_static_function(attr_func)
                if inner_input_spec:
                    inner_input_spec = paddle.utils.pack_sequence_as(
                        input_spec, inner_input_spec
                    )
                static_function = to_static(
                    static_func,
                    input_spec=inner_input_spec,
                    full_graph=True,
                    backend=None,
                )
                concrete_program = static_function.concrete_program

                if static_function.class_instance is None:
                    warnings.warn(
                        f'`jit.save` will only save the `Program`, not the parameters. If you have to save the parameters, please make sure that {layer} is a member function of `paddle.nn.Layer` and the saved parameters are in `state_dict`'
                    )

        # when save multi `StaticFunction`, all `StaticFunction` share params.
        dygraph_state_dict = None
        if isinstance(inner_layer, Layer):
            dygraph_state_dict = inner_layer.to_static_state_dict()
        elif isinstance(attr_func, StaticFunction):
            if static_func.class_instance:
                dygraph_state_dict = (
                    static_func.class_instance.to_static_state_dict()
                )

        if dygraph_state_dict:
            # NOTE(chenweihang): we maintain the mapping of variable name to
            # structured name, the buffer variable (non-persistable)
            # saved to inference program may not need by dygraph Layer,
            # we only record the state_dict variable's structured name
            state_names_dict = {}
            state_var_dict = {}
            for structured_name, var in dygraph_state_dict.items():
                state_names_dict[var.name] = structured_name
                state_var_dict[var.name] = var
        # 3. share parameters from Layer to scope & record var info
        with dygraph.guard():
            if use_pir_api():
                for tensor, value in zip(*concrete_program.parameters):
                    if not value.persistable:
                        continue
                    param_or_buffer_tensor = scope.var(value.name).get_tensor()
                    src_tensor = (
                        state_var_dict[tensor.name].value().get_tensor()
                    )
                    param_or_buffer_tensor._share_data_with(src_tensor)

            else:
                for param_or_buffer in concrete_program.parameters:
                    # share to scope
                    if param_or_buffer.type == core.VarDesc.VarType.VOCAB:
                        scr_tensor = param_or_buffer.value().get_map_tensor()
                        tgt_var = scope.var(param_or_buffer.name)
                        tgt_var.set_vocab(scr_tensor)
                    else:
                        param_or_buffer_tensor = scope.var(
                            param_or_buffer.name
                        ).get_tensor()
                        # src_tensor = param_or_buffer.value().get_tensor()
                        src_tensor = (
                            state_var_dict[param_or_buffer.name]
                            .value()
                            .get_tensor()
                        )
                        param_or_buffer_tensor._share_data_with(src_tensor)
                    # record var info
                    if param_or_buffer.name not in extra_var_info:
                        extra_info_dict = {}
                        if param_or_buffer.name in state_names_dict:
                            extra_info_dict['structured_name'] = (
                                state_names_dict[param_or_buffer.name]
                            )
                        extra_info_dict['stop_gradient'] = (
                            param_or_buffer.stop_gradient
                        )
                        if isinstance(param_or_buffer, EagerParamBase):
                            extra_info_dict['trainable'] = (
                                param_or_buffer.trainable
                            )
                        extra_var_info[param_or_buffer.name] = extra_info_dict
        # 4. build input & output of save_inference_model
        # NOTE(chenweihang): [ Get input variables name ]
        # There are two cases, whether to prune the inputs or not
        # - not prune inputs (recommend):
        #   - the len(input_spec) == len((concrete_program.inputs) - 1
        #   - here can use concrete_program.inputs directly
        # - prune inputs:
        #   - the input_spec length < len((concrete_program.inputs) - 1
        #   - the input_spec's name should be in concrete_program.inputs

        input_vars, input_var_names = _get_input_var_and_names(
            concrete_program.inputs,
            inner_input_spec,
            configs.input_names_after_prune,
        )

        # NOTE(chenweihang): [ Get output variables ]
        # the rule is like [ Get input variables name ]. For output var,
        # we only support Tensor spec, and actually, we only need the
        # var name of output, and we don't recommended to use output_spec
        # NOTE(Ruting): in pir mode, Tensor is not supported, because value has no name in most cases,
        # which can't be used to judge which tensor corresponds to which value; the value can't be found
        # if the saved program is not the same as the program that includes output_spec, so we need to
        # use the position of the output.

        output_vars = _get_output_vars(
            concrete_program.outputs, configs.output_spec, with_hook
        )

        # 5. save inference model
        # construct new save_inference_model arguments
        model_path = dirname
        # NOTE(chenweihang): because prefix contains model and params filename,
        # so we don't support set model_filename & params_filename
        if 'forward' == attr_func or not isinstance(layer, Layer):
            model_filename = file_prefix + INFER_MODEL_SUFFIX
            params_filename = file_prefix + INFER_PARAMS_SUFFIX
            path_prefix = file_prefix
        else:
            model_filename = file_prefix + '.' + attr_func + INFER_MODEL_SUFFIX
            params_filename = (
                file_prefix + '.' + attr_func + INFER_PARAMS_SUFFIX
            )
            path_prefix = file_prefix + '.' + attr_func
        file_path = os.path.join(model_path, path_prefix)
        with scope_guard(scope):
            if use_pir_api():
                value_map = paddle.pir.IrMapping()
                clone_program = concrete_program.main_program.clone(value_map)
                clone_input_vars = []
                for v in input_vars:
                    if type(v) is paddle.static.InputSpec:
                        name = v.name
                        for op in clone_program.global_block().ops:
                            if (
                                op.name() == 'pd_op.data'
                                and op.attrs()["name"] == name
                            ):
                                clone_input_vars.append(op.result(0))
                    else:
                        clone_input_vars.append(value_map.look_up(v))

                clone_output_vars = [value_map.look_up(v) for v in output_vars]

            else:
                input_vars = [
                    concrete_program.main_program.global_block().var(name)
                    for name in input_var_names
                ]
                clone_program = concrete_program.main_program.clone()
                clone_input_vars = input_vars
                clone_output_vars = output_vars
            save_inference_model(
                path_prefix=file_path,
                feed_vars=clone_input_vars,
                fetch_vars=clone_output_vars,
                executor=Executor(_current_expected_place()),
                program=clone_program,
                clip_extra=configs.clip_extra,
                skip_prune_program=configs.skip_prune_program,
                separate_parameters=configs.separate_parameters,
            )

        if combine_params:
            if use_pir_api():
                # NOTE(Ruting): concrete_program has been pruned when init partialProgramLayer,
                # so we do not need to prune again.

                for var in concrete_program.main_program.list_vars():
                    if var.persistable:
                        combine_vars[var.name] = var
                # NOTE(Ruting): concrete_program will delete after this loop item,
                # value delete at the same time, so we use list to Extend its lifecycle
                combine_program.append(concrete_program.main_program)
            else:
                clone_main_program = concrete_program.main_program.clone()
                clone_main_program = clone_main_program._prune_with_input(
                    input_var_names, output_vars
                )
                for block in clone_main_program.blocks:
                    combine_vars.update(block.vars)

    # save shared params
    if combine_params:
        # sort vars by name
        combine_vars = sorted(combine_vars.items(), key=lambda item: item[0])
        ordered_vars = []
        for name, var in combine_vars:
            ordered_vars.append(var)

        params_filename = file_prefix + INFER_PARAMS_SUFFIX
        with scope_guard(scope):
            if use_pir_api():
                paddle.static.save_vars(
                    Executor(_current_expected_place()),
                    dirname=model_path,
                    vars=ordered_vars,
                    filename=params_filename,
                )
            else:
                paddle.static.save_vars(
                    Executor(_current_expected_place()),
                    dirname=model_path,
                    vars=list(
                        filter(
                            paddle.framework.io_utils.is_persistable,
                            ordered_vars,
                        )
                    ),
                    filename=params_filename,
                )
        # save property
        property_save_path = os.path.join(
            os.path.normpath(model_path), file_prefix + INFER_PROPERTY_SUFFIX
        )
        _save_property(property_save_path, property_vals)

    # NOTE(chenweihang): [ Save extra variable info ]
    # save_inference_model will lose some important variable information, including:
    #   - Variable name and correspondence (when saved variables as one file)
    #   - Variable.stop_gradient information
    #   - Which persistent variable are parameter and which are not
    #   - Parameter.trainable information
    #
    # The lost information cannot be recovered when it is loaded again,
    # so if we want to perform fine-tune after loading, we may need to
    # configure redundant information to proceed.
    #
    # Due to compatibility issues, we cannot change the original storage structure,
    # but we can save these information in `jit.save` without changing the original
    # storage to improve user experience. So we save extra information into
    # file `***.pdiparams.info`

    # "layer" can only be Layer or function or StaticFunction.
    contain_parameter = False
    if concrete_program is not None:
        for var in concrete_program.main_program.list_vars():
            if use_pir_api():
                is_persistable = (
                    var.get_defining_op().has_attr("persistable")
                    and var.get_defining_op().attrs()["persistable"] is True
                )
                contain_parameter |= is_persistable
            else:
                contain_parameter |= isinstance(var, Parameter)

    if (isinstance(layer, Layer) or contain_parameter) and extra_var_info:
        with scope_guard(scope):
            extra_var_info_path = path + INFER_PARAMS_INFO_SUFFIX
            with open(extra_var_info_path, 'wb') as f:
                pickle.dump(extra_var_info, f, protocol=2)
    scope.erase(scope.local_var_names())


@dygraph_only
def load(
    path: str, **configs: Unpack[_LoadOptions]
) -> TranslatedLayer | PirTranslatedLayer:
    """
    :api_attr: imperative

    Load model saved by ``paddle.jit.save`` or ``paddle.static.save_inference_model`` or
    paddle 1.x API ``paddle.static.save_inference_model`` as ``paddle.jit.TranslatedLayer``,
    then performing inference or fine-tune training.

    .. note::
        If you load model saved by ``paddle.static.save_inference_model`` ,
        there will be the following limitations when using it in fine-tuning:
        1. Imperative mode do not support DenseTensor. All original model's feed targets or parameters that depend on LoD are temporarily unavailable.
        2. All saved model's feed targets need to be passed into TranslatedLayer's forward function.
        3. The variable's ``stop_gradient`` information is lost and can not be recovered.
        4. The parameter's ``trainable`` information is lost and can not be recovered.

    Args:
        path (str): The path prefix to load model. The format is ``dirname/file_prefix`` or ``file_prefix`` .
        **configs (dict, optional): Other load configuration options for compatibility. We do not
            recommend using these configurations, they may be removed in the future. If not necessary,
            DO NOT use them. Default None.
            The following options are currently supported:
            (1) model_filename (str): The inference model file name of the paddle 1.x
            ``save_inference_model`` save format. Default file name is :code:`__model__` .
            (2) params_filename (str): The persistable variables file name of the paddle 1.x
            ``save_inference_model`` save format. No default file name, save variables separately
            by default.


    Returns:
        TranslatedLayer: A Layer object can run saved translated model.

    Examples:
        1. Load model saved by ``paddle.jit.save`` then performing inference and fine-tune training.

            .. code-block:: pycon
                :name: code-example1

                >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
                >>> import numpy as np
                >>> import paddle
                >>> import paddle.nn as nn
                >>> import paddle.optimizer as opt

                >>> BATCH_SIZE = 16
                >>> BATCH_NUM = 4
                >>> EPOCH_NUM = 4

                >>> IMAGE_SIZE = 784
                >>> CLASS_NUM = 10

                >>> # define a random dataset
                >>> class RandomDataset(paddle.io.Dataset):  # type: ignore[type-arg]
                ...     def __init__(self, num_samples):
                ...         self.num_samples = num_samples
                ...
                ...     def __getitem__(self, idx):
                ...         image = np.random.random([IMAGE_SIZE]).astype('float32')
                ...         label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
                ...         return image, label
                ...
                ...     def __len__(self):
                ...         return self.num_samples

                >>> class LinearNet(nn.Layer):
                ...     def __init__(self):
                ...         super().__init__()
                ...         self._linear = nn.Linear(IMAGE_SIZE, CLASS_NUM)
                ...
                ...     @paddle.jit.to_static
                ...     def forward(self, x):
                ...         return self._linear(x)
                >>> def train(layer, loader, loss_fn, opt):
                ...     for epoch_id in range(EPOCH_NUM):
                ...         for batch_id, (image, label) in enumerate(loader()):
                ...             out = layer(image)
                ...             loss = loss_fn(out, label)
                ...             loss.backward()
                ...             opt.step()
                ...             opt.clear_grad()
                ...             print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))

                >>> # 1. train & save model.

                >>> # create network
                >>> layer = LinearNet()
                >>> loss_fn = nn.CrossEntropyLoss()
                >>> adam = opt.Adam(learning_rate=0.001, parameters=layer.parameters())

                >>> # create data loader
                >>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
                >>> loader = paddle.io.DataLoader(
                ...     dataset,
                ...     batch_size=BATCH_SIZE,
                ...     shuffle=True,
                ...     drop_last=True,
                ...     num_workers=2,
                ... )

                >>> # train
                >>> train(layer, loader, loss_fn, adam)

                >>> # save
                >>> path = "example_model/linear"
                >>> paddle.jit.save(layer, path)

                >>> # 2. load model

                >>> # load
                >>> loaded_layer = paddle.jit.load(path)

                >>> # inference
                >>> loaded_layer.eval()
                >>> x = paddle.randn([1, IMAGE_SIZE], 'float32')
                >>> pred = loaded_layer(x)

                >>> # fine-tune
                >>> loaded_layer.train()
                >>> adam = opt.Adam(learning_rate=0.001, parameters=loaded_layer.parameters())
                >>> train(loaded_layer, loader, loss_fn, adam)


        2. Load model saved by ``paddle.static.save_inference_model`` then performing and fine-tune training.

            .. code-block:: pycon
                :name: code-example2

                >>> # doctest: +SOLO('can not use multiprocessing testing `DataLoader`')
                >>> import numpy as np
                >>> import paddle
                >>> import paddle.static as static
                >>> import paddle.nn as nn
                >>> import paddle.optimizer as opt
                >>> import paddle.nn.functional as F

                >>> BATCH_SIZE = 16
                >>> BATCH_NUM = 4
                >>> EPOCH_NUM = 4

                >>> IMAGE_SIZE = 784
                >>> CLASS_NUM = 10

                >>> # define a random dataset
                >>> class RandomDataset(paddle.io.Dataset):  # type: ignore[type-arg]
                ...     def __init__(self, num_samples):
                ...         self.num_samples = num_samples
                ...
                ...     def __getitem__(self, idx):
                ...         image = np.random.random([IMAGE_SIZE]).astype('float32')
                ...         label = np.random.randint(0, CLASS_NUM - 1, (1,)).astype('int64')
                ...         return image, label
                ...
                ...     def __len__(self):
                ...         return self.num_samples

                >>> paddle.enable_static()

                >>> image = static.data(name='image', shape=[None, 784], dtype='float32')
                >>> label = static.data(name='label', shape=[None, 1], dtype='int64')
                >>> pred = static.nn.fc(x=image, size=10, activation='softmax')
                >>> loss = F.cross_entropy(input=pred, label=label)
                >>> avg_loss = paddle.mean(loss)

                >>> optimizer = paddle.optimizer.SGD(learning_rate=0.001)
                >>> optimizer.minimize(avg_loss)

                >>> place = paddle.CPUPlace()
                >>> exe = static.Executor(place)
                >>> exe.run(static.default_startup_program())

                >>> # create data loader
                >>> dataset = RandomDataset(BATCH_NUM * BATCH_SIZE)
                >>> loader = paddle.io.DataLoader(
                ...     dataset,
                ...     feed_list=[image, label],
                ...     places=place,
                ...     batch_size=BATCH_SIZE,
                ...     shuffle=True,
                ...     drop_last=True,
                ...     return_list=False,
                ...     num_workers=2,
                ... )

                >>> # 1. train and save inference model
                >>> for data in loader():
                ...     exe.run(
                ...         static.default_main_program(),
                ...         feed=data,
                ...         fetch_list=[avg_loss],
                ...     )

                >>> model_path = "fc.example.model"
                >>> paddle.static.save_inference_model(
                ...     model_path,
                ...     [image],
                ...     [pred],
                ...     exe,
                ... )

                >>> # 2. load model

                >>> # enable dygraph mode
                >>> paddle.disable_static(place)

                >>> # load
                >>> fc = paddle.jit.load(model_path)

                >>> # inference
                >>> fc.eval()
                >>> x = paddle.randn([1, IMAGE_SIZE], 'float32')
                >>> pred = fc(x)

                >>> # fine-tune
                >>> fc.train()
                >>> loss_fn = nn.CrossEntropyLoss()
                >>> adam = opt.Adam(learning_rate=0.001, parameters=fc.parameters())
                >>> loader = paddle.io.DataLoader(
                ...     dataset,
                ...     places=place,
                ...     batch_size=BATCH_SIZE,
                ...     shuffle=True,
                ...     drop_last=True,
                ...     num_workers=2,
                ... )
                >>> for epoch_id in range(EPOCH_NUM):
                ...     for batch_id, (image, label) in enumerate(loader()):
                ...         out = fc(image)
                ...         loss = loss_fn(out, label)
                ...         loss.backward()
                ...         adam.step()
                ...         adam.clear_grad()
                ...         print("Epoch {} batch {}: loss = {}".format(epoch_id, batch_id, np.mean(loss.numpy())))
    """
    # 1. construct correct config
    config = _parse_load_config(configs)
    model_path, config = _build_load_path_and_config(path, config)

    if use_pir_api():
        return PirTranslatedLayer._construct(model_path, config)
    else:
        return TranslatedLayer._construct(model_path, config)


def set_dynamic_shape(variable, shape_list):
    if paddle.base.dygraph.base.in_to_static_mode():
        if isinstance(variable, paddle.base.framework.Variable):
            variable.desc.set_shape(shape_list)
        elif isinstance(variable, paddle.pir.Value):
            variable.set_shape(shape_list)
        else:
            raise TypeError(
                "In to_static mode, variable must be a Variable or Value"
            )
    else:
        # in dygraph mode, dynamic shape is not needed, just do nothing.
        return


def get_ast_static_function(function):
    if isinstance(function, SymbolicStaticFunction):
        if function.class_instance:
            dygraph_function = types.MethodType(
                function._dygraph_function, function.class_instance
            )
        else:
            dygraph_function = function._dygraph_function

        if function._function_spec._input_spec is None:
            ast_static_function = ASTStaticFunction(
                dygraph_function,
                function.last_call_input_spec,
                **function._kwargs,
            )
            return ast_static_function
        else:
            ast_static_function = ASTStaticFunction(
                dygraph_function,
                function._function_spec._input_spec,
                **function._kwargs,
            )
            return ast_static_function
    return function


def json_to_pdmodel(net, input_spec, load_path, save_path):
    net1 = paddle.jit.load(load_path)
    state_dict = {}
    for val in net1.state_dict().values():
        name = val.name[: val.name.rfind('_')]
        state_dict[name] = val

    name_state_dict = {}
    for name, var in net.to_static_state_dict().items():
        name_state_dict[name] = state_dict[var.name]

    net.set_state_dict(name_state_dict)
    with paddle.pir_utils.OldIrGuard():
        paddle.jit.save(net, save_path, input_spec)