chore: import upstream snapshot with attribution

python/paddle/incubate/tensor/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2021 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 .manipulation import _npu_identity  # noqa: F401
+from .math import (  # noqa: F401
+    segment_max,
+    segment_mean,
+    segment_min,
+    segment_sum,
+)
+
+__all__ = []

python/paddle/incubate/tensor/manipulation.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2022 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.
+
+import copy
+
+import paddle
+from paddle import _C_ops
+from paddle.base import core
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.framework import EagerParamBase
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+
+__all__ = []
+
+
+# TODO(qili93): remove this op after custom op and custom device
+# integrated and then move this op along with its code to plugin.
+def _npu_identity(x, format=-1):
+    """
+
+    This OP takes in the Tensor :attr:`x` and change it to output with
+    aclFormat with int value. This API is only used for Ascend NPU.
+
+    Args:
+        x(Tensor): An input N-D Tensor with data type bool, float16,
+                   float32, float64, int32, int64, int16, int8, uint8.
+        format(int): Storage data format of the output in aclFormat,
+                     default value is -1.
+
+    Returns:
+        Tensor: A Tensor with acl storage format on Ascend NPU.
+
+    Examples:
+        .. code-block:: pycon
+
+            >>> # doctest: +REQUIRES(env:NPU)
+            >>> import paddle
+            >>> paddle.device.set_device('npu')
+
+            >>> x = paddle.ones(shape=[6])
+            >>> y = paddle.incubate._npu_identity(x, 3)  # ACL_FORMAT_NC1HWC0 = 3
+            >>> print(y.shape)
+            paddle.Size([1, 1, 1, 1, 16])
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.npu_identity(x, format)
+    else:
+        check_variable_and_dtype(
+            x,
+            'x',
+            [
+                'bool',
+                'int8',
+                'uint8',
+                'int16',
+                'int32',
+                'int64',
+                'float16',
+                'float32',
+                'float64',
+            ],
+            'npu_identity',
+        )
+
+        helper = LayerHelper('npu_identity', **locals())
+        out = helper.create_variable_for_type_inference(
+            dtype=x.dtype, stop_gradient=x.stop_gradient
+        )
+        helper.append_op(
+            type='npu_identity',
+            inputs={'x': [x]},
+            outputs={'out': [out]},
+            attrs={'format': format},
+        )
+        return out
+
+
+def _load_reload_impl(src_tensor, func):
+    """
+    Helper to create a new destination tensor and call 'func(dst, src)'
+    which is either offload or reload.
+    """
+    if isinstance(src_tensor, EagerParamBase):
+        state = copy.deepcopy(src_tensor.__dict__)
+        new_param = EagerParamBase(src_tensor.shape, src_tensor.dtype, **state)
+        task = func(new_param, src_tensor)
+        return new_param, task
+    elif isinstance(src_tensor, paddle.Tensor):
+        new_varbase = core.eager.Tensor()
+        task = func(new_varbase, src_tensor)
+        return new_varbase, task
+
+
+def create_async_load():
+    """
+    Constructs a new AsyncLoad object.
+    It is used to load/reload data asynchronously on GPU.
+    """
+    custom_devices = paddle.device.get_all_custom_device_type()
+    if paddle.is_compiled_with_xpu():
+        return core.XpuAsyncLoad()
+    elif any(
+        paddle.is_compiled_with_custom_device(dev) for dev in custom_devices
+    ):
+        return None
+    else:  # default is GPU or CUDA
+        return core.AsyncLoad()
+
+
+def create_xpu_async_load():
+    """
+    Constructs a new AsyncLoad object.
+    It is used to load/reload data asynchronously on XPU.
+    """
+    return core.XpuAsyncLoad()
+
+
+class _NoopAsyncTask:
+    """A dummy Task for sync‐fallback on XPU."""
+
+    def is_completed(self):
+        return True
+
+    def cpu_wait(self):
+        pass
+
+    def xpu_wait(self):
+        pass
+
+
+def async_offload(src_tensor, async_load):
+    """
+    Loads the source tensor into the destination tensor asynchronously.
+
+    Args:
+        src_tensor (EagerParamBase|paddle.Tensor): The source tensor.
+        async_load (core.AsyncLoad): The AsyncLoad object.
+
+    Returns:
+        tuple: A tuple containing two elements:
+         - dest_tensor (EagerParamBase|paddle.Tensor): The destination tensor.
+         - task (Task): The task that loads the source tensor into the destination tensor.
+    """
+    is_xpu_tensor = (
+        paddle.is_compiled_with_xpu()
+        and hasattr(src_tensor, "place")
+        and src_tensor.place.is_xpu_place()
+    )
+
+    # async_offload does not support custom device now
+    custom_devices = paddle.device.get_all_custom_device_type()
+    is_custom_tensor = (
+        any(
+            paddle.is_compiled_with_custom_device(dev) for dev in custom_devices
+        )
+        and hasattr(src_tensor, "place")
+        and src_tensor.place.custom_device_type() in custom_devices
+    )
+
+    if is_xpu_tensor or is_custom_tensor:
+        # sync fallback
+        host_tensor = src_tensor.cpu()
+        out = paddle.to_tensor(host_tensor.numpy(), place=paddle.CPUPlace())
+        return out, _NoopAsyncTask()
+
+    return _load_reload_impl(src_tensor, async_load.offload)
+
+
+def async_reload(src_tensor, async_load):
+    """
+    Reloads the source tensor into the destination tensor asynchronously.
+
+    Args:
+        src_tensor (EagerParamBase|paddle.Tensor): The source tensor.
+        async_load (core.AsyncLoad): The AsyncLoad object.
+
+    Returns:
+        tuple: A tuple containing two elements:
+         - dest_tensor (EagerParamBase|paddle.Tensor): The destination tensor.
+         - task (Task): The task that reloads the source tensor into the destination tensor.
+    """
+
+    if (
+        paddle.is_compiled_with_xpu()
+        and hasattr(src_tensor, "place")
+        and src_tensor.place.is_cpu_place()
+    ):
+        arr = src_tensor.numpy()
+        xpu = paddle.to_tensor(arr, place=paddle.XPUPlace(0))
+        return xpu, _NoopAsyncTask()
+
+    return _load_reload_impl(src_tensor, async_load.reload)
+
+
+def async_offload_with_offset(
+    src_tensor, dst_tensor, src_offset, dst_offset, offload_size, async_loader
+):
+    """
+    Offloading the source tensor into the destination tensor asynchronously with offset and size customized.
+
+    Args:
+        src_tensor (EagerParamBase|paddle.Tensor): The source tensor.
+        dst_tensor (EagerParamBase|paddle.Tensor): The destination tensor.
+        src_offset (int): The element offset of the source tensor.
+        dst_offset (int): The element offset of the destination tensor.
+        offload_size (int): The size of the data to be loaded.
+        async_loader (core.AsyncLoad): The AsyncLoad object.
+
+    Returns:
+        task (Task): The task that operates partial offloading.
+    """
+    assert len(src_tensor.shape) <= 1, "Only support 1-D tensor"
+    assert len(dst_tensor.shape) <= 1, "Only support 1-D tensor"
+    assert src_tensor.dtype == dst_tensor.dtype, "Only support same dtype"
+
+    return async_loader.offload_with_offset(
+        dst_tensor, src_tensor, dst_offset, src_offset, offload_size
+    )
+
+
+def enable_activation_offload(model, enable=True, retry_times=1):
+    """
+    Enable activation offload
+    """
+    if enable:
+        paddle.set_flags({"FLAGS_offload_retry_times": retry_times})
+        paddle.core.register_offload_callback()
+        paddle.core.set_skip_offload_callback_tensors(model.parameters())
+    else:
+        paddle.set_flags({"FLAGS_offload_retry_times": -1})
+        paddle.core.clear_offload_callback()
+        paddle.core.set_skip_offload_callback_tensors([])

python/paddle/incubate/tensor/math.py

@@ -0,0 +1,309 @@
+# Copyright (c) 2021 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
+
+from paddle import _C_ops
+from paddle.base.data_feeder import check_variable_and_dtype
+from paddle.base.layer_helper import LayerHelper
+from paddle.framework import in_dynamic_or_pir_mode
+from paddle.utils import deprecated
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+__all__ = []
+
+
+@deprecated(
+    since="2.4.0",
+    update_to="paddle.geometric.segment_sum",
+    level=1,
+    reason="paddle.incubate.segment_sum will be removed in future",
+)
+def segment_sum(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment Sum Operator.
+
+    Sum the elements of input `data` which with
+    the same index in `segment_ids`.
+    It computes a tensor such that
+
+    .. math::
+
+        out_i = \sum_{j \in \{segment\_ids_j == i \} } data_{j}
+
+    where sum is over j such that `segment_ids[j] == i`.
+
+    Args:
+        data (Tensor): A tensor, available data type float32, float64, int32, int64.
+        segment_ids (Tensor): A 1-D tensor, which have the same size
+                            with the first dimension of input data.
+                            Available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+       Tensor, the Segment Sum result.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.incubate.segment_sum(data, segment_ids)
+            >>> print(out)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[4., 4., 4.],
+             [4., 5., 6.]])
+
+    """
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "SUM")
+    else:
+        check_variable_and_dtype(
+            data, "X", ("float32", "float64", "int32", "int64"), "segment_pool"
+        )
+        check_variable_and_dtype(
+            segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+        )
+
+        helper = LayerHelper("segment_sum", **locals())
+        out = helper.create_variable_for_type_inference(dtype=data.dtype)
+        summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+        helper.append_op(
+            type="segment_pool",
+            inputs={"X": data, "SegmentIds": segment_ids},
+            outputs={"Out": out, "SummedIds": summed_ids},
+            attrs={"pooltype": "SUM"},
+        )
+        return out
+
+
+@deprecated(
+    since="2.4.0",
+    update_to="paddle.geometric.segment_mean",
+    level=1,
+    reason="paddle.incubate.segment_mean will be removed in future",
+)
+def segment_mean(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment Mean Operator.
+
+    Ihis operator calculate the mean value of input `data` which
+    with the same index in `segment_ids`.
+    It computes a tensor such that
+
+    .. math::
+
+        out_i = \mathop{mean}_{j \in \{segment\_ids_j == i \} } data_{j}
+
+    where sum is over j such that 'segment_ids[j] == i' and $n_i$ is the number
+    of all index 'segment_ids[j] == i'.
+
+    Args:
+        data (tensor): a tensor, available data type float32, float64, int32, int64.
+        segment_ids (tensor): a 1-d tensor, which have the same size
+                            with the first dimension of input data.
+                            available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+       Tensor, the Segment Mean result.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.incubate.segment_mean(data, segment_ids)
+            >>> print(out)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[2., 2., 2.],
+             [4., 5., 6.]])
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "MEAN")
+
+    check_variable_and_dtype(
+        data, "X", ("float32", "float64", "int32", "int64"), "segment_pool"
+    )
+    check_variable_and_dtype(
+        segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+    )
+
+    helper = LayerHelper("segment_mean", **locals())
+    out = helper.create_variable_for_type_inference(dtype=data.dtype)
+    summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+    helper.append_op(
+        type="segment_pool",
+        inputs={"X": data, "SegmentIds": segment_ids},
+        outputs={"Out": out, "SummedIds": summed_ids},
+        attrs={"pooltype": "MEAN"},
+    )
+    return out
+
+
+@deprecated(
+    since="2.4.0",
+    update_to="paddle.geometric.segment_min",
+    level=1,
+    reason="paddle.incubate.segment_min will be removed in future",
+)
+def segment_min(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment min operator.
+
+    Calculate the minimum elements of input `data` which with
+    the same index in `segment_ids`.
+    It computes a tensor such that
+
+    .. math::
+
+        out_i = \min_{j \in \{segment\_ids_j == i \} } data_{j}
+
+    where min is over j such that `segment_ids[j] == i`.
+
+    Args:
+        data (tensor): a tensor, available data type float32, float64, int32, int64.
+        segment_ids (tensor): a 1-d tensor, which have the same size
+                            with the first dimension of input data.
+                            available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+       Tensor, the minimum result.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.incubate.segment_min(data, segment_ids)
+            >>> print(out)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[1., 2., 1.],
+             [4., 5., 6.]])
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        return _C_ops.segment_pool(data, segment_ids, "MIN")
+
+    check_variable_and_dtype(
+        data, "X", ("float32", "float64", "int32", "int64"), "segment_pool"
+    )
+    check_variable_and_dtype(
+        segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+    )
+
+    helper = LayerHelper("segment_min", **locals())
+    out = helper.create_variable_for_type_inference(dtype=data.dtype)
+    summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+    helper.append_op(
+        type="segment_pool",
+        inputs={"X": data, "SegmentIds": segment_ids},
+        outputs={"Out": out, "SummedIds": summed_ids},
+        attrs={"pooltype": "MIN"},
+    )
+    return out
+
+
+@deprecated(
+    since="2.4.0",
+    update_to="paddle.geometric.segment_max",
+    level=1,
+    reason="paddle.incubate.segment_max will be removed in future",
+)
+def segment_max(
+    data: Tensor, segment_ids: Tensor, name: str | None = None
+) -> Tensor:
+    r"""
+    Segment max operator.
+
+    Calculate the maximum elements of input `data` which with
+    the same index in `segment_ids`.
+    It computes a tensor such that
+
+    .. math::
+
+        out_i = \max_{j \in \{segment\_ids_j == i \} } data_{j}
+
+    where max is over j such that `segment_ids[j] == i`.
+
+    Args:
+        data (tensor): a tensor, available data type float32, float64, int32, int64.
+        segment_ids (tensor): a 1-d tensor, which have the same size
+                            with the first dimension of input data.
+                            available data type is int32, int64.
+        name (str, optional): Name for the operation (optional, default is None).
+                            For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+       Tensor, the maximum result.
+
+    Examples:
+
+        .. code-block:: pycon
+
+            >>> import paddle
+            >>> data = paddle.to_tensor([[1, 2, 3], [3, 2, 1], [4, 5, 6]], dtype='float32')
+            >>> segment_ids = paddle.to_tensor([0, 0, 1], dtype='int32')
+            >>> out = paddle.incubate.segment_max(data, segment_ids)
+            >>> print(out)
+            Tensor(shape=[2, 3], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[3., 2., 3.],
+             [4., 5., 6.]])
+
+    """
+
+    if in_dynamic_or_pir_mode():
+        out = _C_ops.segment_pool(data, segment_ids, "MAX")
+        return out
+
+    check_variable_and_dtype(
+        data, "X", ("float32", "float64", "int32", "int64"), "segment_pool"
+    )
+    check_variable_and_dtype(
+        segment_ids, "SegmentIds", ("int32", "int64"), "segment_pool"
+    )
+
+    helper = LayerHelper("segment_max", **locals())
+    out = helper.create_variable_for_type_inference(dtype=data.dtype)
+    summed_ids = helper.create_variable_for_type_inference(dtype=data.dtype)
+    helper.append_op(
+        type="segment_pool",
+        inputs={"X": data, "SegmentIds": segment_ids},
+        outputs={"Out": out, "SummedIds": summed_ids},
+        attrs={"pooltype": "MAX"},
+    )
+    return out