chore: import upstream snapshot with attribution

python/paddle/incubate/framework/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2023 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 .random import (  # noqa: F401
+    get_rng_state,
+    register_rng_state_as_index,
+    set_rng_state,
+)
+
+__all__ = []

python/paddle/incubate/framework/random.py

@@ -0,0 +1,311 @@
+#   Copyright (c) 2023 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
+
+from typing import TYPE_CHECKING, overload
+
+import paddle
+from paddle import base
+from paddle.base import core
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from typing import Literal, Protocol
+
+    from paddle._typing import PlaceLike
+
+    class _GeneratorState(Protocol):
+        def current_seed(self) -> int: ...
+
+
+__all__ = []
+
+
+@overload
+def get_rng_state(
+    device: PlaceLike | None = ..., use_index: Literal[True] = ...
+) -> list[int]: ...
+
+
+@overload
+def get_rng_state(
+    device: PlaceLike | None = ..., use_index: Literal[False] = ...
+) -> list[_GeneratorState]: ...
+
+
+@overload
+def get_rng_state(
+    device: PlaceLike | None = ..., use_index: bool = ...
+) -> list[int] | list[_GeneratorState]: ...
+
+
+def get_rng_state(device=None, use_index=False):
+    """
+    Get all random states of random generators of specified device.
+
+    Args:
+        device(str): This parameter determines the specific running device.
+            It can be ``cpu``, ``gpu``, ``xpu``, Default is None.
+            If None, return the generators of current device (specified by ``set_device``).
+        use_index(bool): If use index is True, return the index that saved in the generator
+
+    Returns:
+        GeneratorState:  object.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> sts = paddle.incubate.get_rng_state()
+    """
+
+    def get_state(generator):
+        if use_index:
+            return generator.get_state_index()
+        else:
+            return generator.get_state()
+
+    state_list = []
+    if device is None:
+        place = base.framework._current_expected_place_()
+    else:
+        place = paddle.device._convert_to_place(device)
+
+    if isinstance(place, core.CPUPlace):
+        state_list.append(get_state(core.default_cpu_generator()))
+    elif isinstance(place, core.CUDAPlace):
+        for i in range(core.get_cuda_device_count()):
+            state_list.append(get_state(core.default_cuda_generator(i)))
+    elif isinstance(place, core.XPUPlace):
+        for i in range(core.get_xpu_device_count()):
+            state_list.append(get_state(core.default_xpu_generator(i)))
+    elif isinstance(place, core.CustomPlace):
+        dev_cnt = sum(
+            [
+                place.get_device_type() == s.split(':')[0]
+                for s in core.get_available_custom_device()
+            ]
+        )
+        for i in range(dev_cnt):
+            state_list.append(
+                get_state(
+                    core.default_custom_device_generator(
+                        core.CustomPlace(place.get_device_type(), i)
+                    )
+                )
+            )
+    else:
+        raise ValueError(
+            f"get_rng_state is not implemented for current device: {place}"
+        )
+
+    return state_list
+
+
+@overload
+def set_rng_state(
+    state_list: Sequence[int],
+    device: PlaceLike | None = ...,
+    use_index: Literal[True] = ...,
+) -> None: ...
+
+
+@overload
+def set_rng_state(
+    state_list: Sequence[_GeneratorState],
+    device: PlaceLike | None = ...,
+    use_index: Literal[False] = ...,
+) -> None: ...
+
+
+@overload
+def set_rng_state(
+    state_list: Sequence[int] | Sequence[_GeneratorState],
+    device: PlaceLike | None = ...,
+    use_index: bool = ...,
+) -> None: ...
+
+
+def set_rng_state(
+    state_list,
+    device=None,
+    use_index=False,
+):
+    """
+
+    Sets generator state for all device generators.
+
+    Args:
+        state_list(list|tuple): The device states to set back to device generators. state_list is obtained from get_rng_state().
+        device(str): This parameter determines the specific running device.
+            It can be ``cpu``, ``gpu``, ``xpu``, Default is None.
+            If None, return the generators of current device (specified by ``set_device``).
+        use_index(bool): If use index is True, state_list should be the indices of the states
+
+    Returns:
+        None.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> sts = paddle.incubate.get_rng_state()
+            >>> paddle.incubate.set_rng_state(sts)
+
+    """
+
+    def set_state(generator, state):
+        if use_index:
+            generator.set_state_index(state)
+        else:
+            generator.set_state(state)
+
+    if device is None:
+        place = base.framework._current_expected_place_()
+    else:
+        place = device._convert_to_place(device)
+
+    if isinstance(place, core.CUDAPlace):
+        if not len(state_list) == core.get_cuda_device_count():
+            raise ValueError(
+                "Length of gpu state list should be equal to the gpu device count"
+            )
+        for i in range(core.get_cuda_device_count()):
+            set_state(core.default_cuda_generator(i), state_list[i])
+    elif isinstance(place, core.XPUPlace):
+        if not len(state_list) == core.get_xpu_device_count():
+            raise ValueError(
+                "Length of xpu state list should be equal to the xpu device count"
+            )
+        for i in range(core.get_xpu_device_count()):
+            set_state(core.default_xpu_generator(i), state_list[i])
+    elif isinstance(place, core.CustomPlace):
+        dev_types = core.get_all_custom_device_type()
+        dev_type = dev_types[0]
+        dev_cnt = core.get_custom_device_count(dev_type)
+        if not len(state_list) == dev_cnt:
+            raise ValueError(
+                f"Length of custom device state list should be equal to the {dev_cnt} device count"
+            )
+        for i in range(dev_cnt):
+            set_state(
+                core.default_custom_device_generator(
+                    core.CustomPlace(place.get_device_type(), i)
+                ),
+                state_list[i],
+            )
+    elif isinstance(place, core.CPUPlace):
+        if not len(state_list) == 1:
+            raise ValueError("Length of cpu state list should be equal to 1")
+        set_state(core.default_cpu_generator(), state_list[0])
+    else:
+        raise ValueError(
+            f"set_rng_state is not implemented for current device: {place}"
+        )
+
+
+def register_rng_state_as_index(
+    state_list: Sequence[_GeneratorState] | None = None,
+    device: PlaceLike | None = None,
+) -> list[int]:
+    """
+
+    The register_rng_state_as_index function creates and registers a new generator state within the generator.
+    It enables users to manage multiple generator states via indices,
+    offering a convenient way to switch between these states without directly manipulating the generator's state.
+
+    Args:
+        state_list(list|tuple): A list or tuple representing the RNG states for devices.
+            If not provided, the function will register the current state.
+        device(str): This parameter determines the specific running device.
+            It can be ``cpu``, ``gpu``, ``xpu``, Default is None.
+            If None, return the generators of current device (specified by ``set_device``).
+
+    Returns:
+        A list of indices representing the positions at which the new states were saved within the generator.
+        These indices can be used to switch between states using set_rng_state(use_index=True)
+
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> old_index = paddle.incubate.get_rng_state(use_index=True)
+            >>> print(old_index)
+            [0]
+            >>> new_index = paddle.incubate.register_rng_state_as_index()
+            >>> print(new_index)
+            [1]
+            >>> paddle.incubate.set_rng_state(old_index, use_index=True)
+            >>> paddle.incubate.set_rng_state(new_index, use_index=True)
+
+    """
+    new_state_index_list = []
+
+    if device is None:
+        place = base.framework._current_expected_place()
+    else:
+        place = device._convert_to_place(device)
+
+    if state_list is None:
+        state_list = get_rng_state(device)
+
+    if isinstance(place, core.CUDAPlace):
+        if not len(state_list) == core.get_cuda_device_count():
+            raise ValueError(
+                "Length of gpu state list should be equal to the gpu device count"
+            )
+        for i in range(core.get_cuda_device_count()):
+            new_state_index_list.append(
+                core.default_cuda_generator(i).register_state_index(
+                    state_list[i]
+                )
+            )
+    elif isinstance(place, core.XPUPlace):
+        if not len(state_list) == core.get_xpu_device_count():
+            raise ValueError(
+                "Length of xpu state list should be equal to the xpu device count"
+            )
+        for i in range(core.get_xpu_device_count()):
+            new_state_index_list.append(
+                core.default_xpu_generator(i).register_state_index(
+                    state_list[i]
+                )
+            )
+    elif isinstance(place, core.CustomPlace):
+        dev_types = core.get_all_custom_device_type()
+        dev_type = dev_types[0]
+        dev_cnt = core.get_custom_device_count(dev_type)
+        if not len(state_list) == dev_cnt:
+            raise ValueError(
+                f"Length of custom device state list should be equal to the {dev_cnt} device count"
+            )
+        for i in range(dev_cnt):
+            new_state_index_list.append(
+                core.default_custom_device_generator(
+                    core.CustomPlace(place.get_device_type(), i)
+                ).register_state_index(state_list[i])
+            )
+    elif isinstance(place, core.CPUPlace):
+        if not len(state_list) == 1:
+            raise ValueError("Length of cpu state list should be equal to 1")
+        new_state_index_list.append(
+            core.default_cpu_generator().register_state_index(state_list[0])
+        )
+    else:
+        raise ValueError(
+            f"register_rng_state_index is not implemented for current device: {place}"
+        )
+    return new_state_index_list