stat.py 29.1 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
#   Copyright (c) 2020 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.

15
# TODO: define statistical functions of a tensor
16

17
import paddle
Z
zyfncg 已提交
18
from paddle import _C_ops
19
from paddle.fluid.framework import in_dygraph_mode
20

21
from ..common_ops_import import Variable
22 23
from ..fluid.data_feeder import check_type, check_variable_and_dtype
from ..framework import LayerHelper, core
24
from .math import _get_reduce_axis_with_tensor
25
from .search import where
26

27 28
__all__ = []

29 30 31 32 33 34

def mean(x, axis=None, keepdim=False, name=None):
    """
    Computes the mean of the input tensor's elements along ``axis``.

    Args:
35
        x (Tensor): The input Tensor with data type float32, float64.
36 37 38 39 40 41 42
        axis (int|list|tuple, optional): The axis along which to perform mean
            calculations. ``axis`` should be int, list(int) or tuple(int). If
            ``axis`` is a list/tuple of dimension(s), mean is calculated along
            all element(s) of ``axis`` . ``axis`` or element(s) of ``axis``
            should be in range [-D, D), where D is the dimensions of ``x`` . If
            ``axis`` or element(s) of ``axis`` is less than 0, it works the
            same way as :math:`axis + D` . If ``axis`` is None, mean is
43
            calculated over all elements of ``x``. Default is None.
44
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
45
            in the output Tensor. If ``keepdim`` is True, the dimensions of
46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is False.
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor, results of average along ``axis`` of ``x``, with the same data
        type as ``x``.

    Examples:
        .. code-block:: python

            import paddle

Z
zhupengyang 已提交
61 62 63 64 65 66
            x = paddle.to_tensor([[[1., 2., 3., 4.],
                                   [5., 6., 7., 8.],
                                   [9., 10., 11., 12.]],
                                  [[13., 14., 15., 16.],
                                   [17., 18., 19., 20.],
                                   [21., 22., 23., 24.]]])
67
            out1 = paddle.mean(x)
68
            # 12.5
69 70 71 72 73 74 75 76 77 78 79 80 81
            out2 = paddle.mean(x, axis=-1)
            # [[ 2.5  6.5 10.5]
            #  [14.5 18.5 22.5]]
            out3 = paddle.mean(x, axis=-1, keepdim=True)
            # [[[ 2.5]
            #   [ 6.5]
            #   [10.5]]
            #  [[14.5]
            #   [18.5]
            #   [22.5]]]
            out4 = paddle.mean(x, axis=[0, 2])
            # [ 8.5 12.5 16.5]
    """
82
    if in_dygraph_mode():
83
        return _C_ops.mean(x, axis, keepdim)
84 85 86 87 88 89 90
    else:
        reduce_all, axis = _get_reduce_axis_with_tensor(axis, x)
        check_variable_and_dtype(
            x,
            'x/input',
            ['uint16', 'float16', 'float32', 'float64'],
            'mean/reduce_mean',
91
        )
92 93 94 95 96 97 98 99 100 101 102
        check_type(
            axis, 'axis/dim', (int, list, tuple, Variable), 'mean/reduce_mean'
        )
        if isinstance(axis, (list, tuple)):
            for item in axis:
                check_type(
                    item,
                    'elements of axis/dim',
                    (int, Variable),
                    'mean/reduce_mean',
                )
103

104
        helper = LayerHelper('mean', **locals())
105

106 107 108 109 110 111 112 113 114
        attrs = {'dim': axis, 'keep_dim': keepdim, 'reduce_all': reduce_all}
        out = helper.create_variable_for_type_inference(x.dtype)
        helper.append_op(
            type='reduce_mean',
            inputs={'X': x},
            outputs={'Out': out},
            attrs=attrs,
        )
        return out
115 116


117
def var(x, axis=None, unbiased=True, keepdim=False, name=None):
118
    """
119
    Computes the variance of ``x`` along ``axis`` .
120 121

    Args:
L
LoneRanger 已提交
122
        x (Tensor): The input Tensor with data type float16, float32, float64.
123 124 125 126
        axis (int|list|tuple, optional): The axis along which to perform variance calculations. ``axis`` should be int, list(int) or tuple(int).

            - If ``axis`` is a list/tuple of dimension(s), variance is calculated along all element(s) of ``axis`` . ``axis`` or element(s) of ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
            - If ``axis`` or element(s) of ``axis`` is less than 0, it works the same way as :math:`axis + D` .
127 128 129 130 131
            - If ``axis`` is None, variance is calculated over all elements of ``x``. Default is None.

        unbiased (bool, optional): Whether to use the unbiased estimation. If ``unbiased`` is True, the divisor used in the computation is :math:`N - 1`, where :math:`N` represents the number of elements along ``axis`` , otherwise the divisor is :math:`N`. Default is True.
        keep_dim (bool, optional): Whether to reserve the reduced dimension in the output Tensor. The result tensor will have one fewer dimension than the input unless keep_dim is true. Default is False.
        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
132 133

    Returns:
134
        Tensor, results of variance along ``axis`` of ``x``, with the same data type as ``x``.
135 136 137 138 139

    Examples:
        .. code-block:: python

            import paddle
140

Z
zhupengyang 已提交
141
            x = paddle.to_tensor([[1.0, 2.0, 3.0], [1.0, 4.0, 5.0]])
142
            out1 = paddle.var(x)
143
            # 2.66666667
144 145
            out2 = paddle.var(x, axis=1)
            # [1.         4.33333333]
146
    """
147
    if not in_dygraph_mode():
L
LoneRanger 已提交
148 149 150
        check_variable_and_dtype(
            x, 'x', ['float16', 'float32', 'float64'], 'var'
        )
151 152

    u = mean(x, axis, True, name)
153
    out = paddle.sum(paddle.pow((x - u), 2), axis, keepdim=keepdim, name=name)
154

155
    dtype = x.dtype
156 157 158
    n = paddle.cast(paddle.numel(x), paddle.int64) / paddle.cast(
        paddle.numel(out), paddle.int64
    )
159
    n = n.astype(dtype)
160
    if unbiased:
161
        one_const = paddle.ones([], x.dtype)
162
        n = where(n > one_const, n - 1.0, one_const)
163
    n.stop_gradient = True
164 165 166
    out /= n
    return out

S
swtkiwi 已提交
167

168 169 170
def std(x, axis=None, unbiased=True, keepdim=False, name=None):
    """
    Computes the standard-deviation of ``x`` along ``axis`` .
L
Liufang Sang 已提交
171 172

    Args:
L
LoneRanger 已提交
173
        x (Tensor): The input Tensor with data type float16, float32, float64.
174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195
        axis (int|list|tuple, optional): The axis along which to perform
            standard-deviation calculations. ``axis`` should be int, list(int)
            or tuple(int). If ``axis`` is a list/tuple of dimension(s),
            standard-deviation is calculated along all element(s) of ``axis`` .
            ``axis`` or element(s) of ``axis`` should be in range [-D, D),
            where D is the dimensions of ``x`` . If ``axis`` or element(s) of
            ``axis`` is less than 0, it works the same way as :math:`axis + D` .
            If ``axis`` is None, standard-deviation is calculated over all
            elements of ``x``. Default is None.
        unbiased (bool, optional): Whether to use the unbiased estimation. If
            ``unbiased`` is True, the standard-deviation is calculated via the
            unbiased estimator. If ``unbiased`` is True,  the divisor used in
            the computation is :math:`N - 1`, where :math:`N` represents the
            number of elements along ``axis`` , otherwise the divisor is
            :math:`N`. Default is True.
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
            in the output Tensor. If ``keepdim`` is True, the dimensions of
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is False.
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.
L
Liufang Sang 已提交
196 197

    Returns:
198 199 200
        Tensor, results of standard-deviation along ``axis`` of ``x``, with the
        same data type as ``x``.

L
Liufang Sang 已提交
201 202 203 204
    Examples:
        .. code-block:: python

            import paddle
205

Z
zhupengyang 已提交
206
            x = paddle.to_tensor([[1.0, 2.0, 3.0], [1.0, 4.0, 5.0]])
207
            out1 = paddle.std(x)
208
            # 1.63299316
209
            out2 = paddle.std(x, unbiased=False)
210
            # 1.49071205
211
            out3 = paddle.std(x, axis=1)
212
            # [1.       2.081666]
213

L
Liufang Sang 已提交
214
    """
215
    if not in_dygraph_mode():
L
LoneRanger 已提交
216 217 218
        check_variable_and_dtype(
            x, 'x', ['float16', 'float32', 'float64'], 'std'
        )
219 220
    out = var(**locals())
    return paddle.sqrt(out)
221 222 223 224


def numel(x, name=None):
    """
225
    Returns the number of elements for a tensor, which is a 0-D int64 Tensor with shape [].
226 227 228

    Args:
        x (Tensor): The input Tensor, it's data type can be bool, float16, float32, float64, int32, int64.
229 230
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.
231 232

    Returns:
233
        Tensor: The number of elements for the input Tensor, whose shape is [].
234 235 236 237

    Examples:
        .. code-block:: python

238
            import paddle
239

240 241
            x = paddle.full(shape=[4, 5, 7], fill_value=0, dtype='int32')
            numel = paddle.numel(x) # 140
242 243 244


    """
245
    if in_dygraph_mode():
246
        return _C_ops.numel(x)
247 248 249 250 251 252 253 254 255
    else:
        if not isinstance(x, Variable):
            raise TypeError("x must be a Tensor in numel")
        helper = LayerHelper('numel', **locals())
        out = helper.create_variable_for_type_inference(
            dtype=core.VarDesc.VarType.INT64
        )
        helper.append_op(type='size', inputs={'Input': x}, outputs={'Out': out})
        return out
Z
zhulei 已提交
256 257


258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320
def nanmedian(x, axis=None, keepdim=True, name=None):
    r"""
    Compute the median along the specified axis, while ignoring NaNs.

    If the valid count of elements is a even number,
    the average value of both elements in the middle is calculated as the median.

    Args:
        x (Tensor): The input Tensor, it's data type can be int32, int64, float16, float32, float64.
        axis (None|int|list|tuple, optional):
            The axis along which to perform median calculations ``axis`` should be int or list of int.
            ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
            If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
            If ``axis`` is None, median is calculated over all elements of ``x``. Default is None.
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
            in the output Tensor. If ``keepdim`` is True, the dimensions of
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is True.
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor, results of median along ``axis`` of ``x``. The output dtype is the same as `x`.

    Examples:
        .. code-block:: python

            import paddle
            x = paddle.to_tensor([[float('nan'), 2. , 3. ], [0. , 1. , 2. ]])

            y1 = x.nanmedian()
            # y1 is [[2.]]

            y2 = x.nanmedian(0)
            # y2 is [[0.,  1.5, 2.5]]

            y3 = x.nanmedian(0, keepdim=False)
            # y3 is [0.,  1.5, 2.5]

            y4 = x.nanmedian((0, 1))
            # y4 is [[2.]]
    """
    if not isinstance(x, Variable):
        raise TypeError("In median, the input x should be a Tensor.")

    if isinstance(axis, (list, tuple)) and len(axis) == 0:
        raise ValueError("Axis list should not be empty.")

    dims = len(x.shape)
    if axis is None:
        axis = []
    elif isinstance(axis, tuple):
        axis = list(axis)
    elif isinstance(axis, int):
        axis = [axis]

    if not isinstance(axis, list):
        raise ValueError(
            "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
        )

    for i in range(len(axis)):
321 322 323
        if not isinstance(axis[i], int) or not (
            axis[i] < dims and axis[i] >= -dims
        ):
324 325 326 327 328 329 330 331 332
            raise ValueError(
                "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
            )
        if axis[i] < 0:
            axis[i] += dims

    if len(axis) != len(set(axis)):
        raise ValueError("Axis has duplicated elements.")

333
    if in_dygraph_mode():
Z
zyfncg 已提交
334
        return _C_ops.nanmedian(x, axis, keepdim)
335 336 337 338 339 340 341
    else:
        check_variable_and_dtype(
            x,
            'X',
            ['int32', 'int64', 'float16', 'float32', 'float64'],
            'nanmedian',
        )
342

343 344 345 346 347 348 349 350 351 352 353
        helper = LayerHelper('nanmedian', **locals())
        attrs = {'axis': axis, 'keepdim': keepdim}
        out = helper.create_variable_for_type_inference(x.dtype)
        medians = helper.create_variable_for_type_inference(x.dtype)
        helper.append_op(
            type='nanmedian',
            inputs={'X': x},
            outputs={'Out': out, 'MedianIndex': medians},
            attrs=attrs,
        )
        return out
354 355


Z
zhulei 已提交
356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382
def median(x, axis=None, keepdim=False, name=None):
    """
    Compute the median along the specified axis.

    Args:
        x (Tensor): The input Tensor, it's data type can be bool, float16, float32, float64, int32, int64.
        axis (int, optional): The axis along which to perform median calculations ``axis`` should be int.
            ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
            If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
            If ``axis`` is None, median is calculated over all elements of ``x``. Default is None.
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
            in the output Tensor. If ``keepdim`` is True, the dimensions of
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is False.
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor, results of median along ``axis`` of ``x``. If data type of ``x`` is float64, data type of results will be float64, otherwise data type will be float32.

    Examples:
        .. code-block:: python

            import paddle

            x = paddle.arange(12).reshape([3, 4])
383 384 385 386
            # Tensor(shape=[3, 4], dtype=int64, place=Place(cpu), stop_gradient=True,
            #        [[0 , 1 , 2 , 3 ],
            #         [4 , 5 , 6 , 7 ],
            #         [8 , 9 , 10, 11]])
Z
zhulei 已提交
387 388

            y1 = paddle.median(x)
389 390
            # Tensor(shape=[], dtype=float32, place=Place(cpu), stop_gradient=True,
            #        5.50000000)
Z
zhulei 已提交
391 392

            y2 = paddle.median(x, axis=0)
393 394
            # Tensor(shape=[4], dtype=float32, place=Place(cpu), stop_gradient=True,
            #        [4., 5., 6., 7.])
Z
zhulei 已提交
395 396

            y3 = paddle.median(x, axis=1)
397 398
            # Tensor(shape=[3], dtype=float32, place=Place(cpu), stop_gradient=True,
            #        [1.50000000, 5.50000000, 9.50000000])
Z
zhulei 已提交
399 400

            y4 = paddle.median(x, axis=0, keepdim=True)
401 402
            # Tensor(shape=[1, 4], dtype=float32, place=Place(cpu), stop_gradient=True,
            #        [[4., 5., 6., 7.]])
Z
zhulei 已提交
403 404 405 406

    """
    if not isinstance(x, Variable):
        raise TypeError("In median, the input x should be a Tensor.")
407

408 409 410
    if x.size == 0:
        raise ValueError("In median, the size of input x should not be 0.")

411
    is_flatten = False
Z
zhulei 已提交
412
    dims = len(x.shape)
413 414 415 416 417
    if dims == 0:
        assert axis in [
            -1,
            0,
            None,
418
        ], 'when input 0-D, axis can only be [-1, 0] or default None'
419 420 421 422 423
        is_flatten = True

    if axis is None:
        is_flatten = True

Z
zhulei 已提交
424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439
    if is_flatten:
        x = paddle.flatten(x)
        axis = 0
    else:
        if not isinstance(axis, int) or not (axis < dims and axis >= -dims):
            raise ValueError(
                "In median, axis should be none or an integer in range [-rank(x), rank(x))."
            )
        if axis < 0:
            axis += dims
    sz = x.shape[axis]
    kth = sz >> 1
    tensor_topk, idx = paddle.topk(x, kth + 1, axis=axis, largest=False)
    dtype = 'float64' if x.dtype == core.VarDesc.VarType.FP64 else 'float32'
    if sz & 1 == 0:
        out_tensor = paddle.slice(
440 441
            tensor_topk, axes=[axis], starts=[kth - 1], ends=[kth]
        ) + paddle.slice(tensor_topk, axes=[axis], starts=[kth], ends=[kth + 1])
Z
zhulei 已提交
442 443
        out_tensor = paddle.cast(out_tensor, dtype=dtype) / 2
    else:
444 445 446 447 448 449
        out_tensor = paddle.cast(
            paddle.slice(
                tensor_topk, axes=[axis], starts=[kth], ends=[kth + 1]
            ),
            dtype=dtype,
        )
450
    out_tensor = out_tensor + paddle.sum(
451 452
        paddle.cast(paddle.isnan(x), dtype=dtype) * x, axis=axis, keepdim=True
    )
453 454 455
    if is_flatten:
        if keepdim:
            out_tensor = out_tensor.reshape([1] * dims)
Z
zhulei 已提交
456
        else:
457
            out_tensor = out_tensor.reshape([])
Z
zhulei 已提交
458
    else:
459 460
        if not keepdim:
            out_tensor = out_tensor.squeeze(axis)
Z
zhulei 已提交
461
    return out_tensor
462 463


464
def _compute_quantile(x, q, axis=None, keepdim=False, ignore_nan=False):
465 466 467
    """
    Compute the quantile of the input along the specified axis.

468
    Args:
469
        x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
470
        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
471 472 473 474 475 476 477 478 479 480 481
            each q will be calculated and the first dimension of output is same to the number of ``q`` .
        axis (int|list, optional): The axis along which to calculate quantile. ``axis`` should be int or list of int.
            ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
            If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
            If ``axis`` is a list, quantile is calculated over all elements of given axises.
            If ``axis`` is None, quantile is calculated over all elements of ``x``. Default is None.
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
            in the output Tensor. If ``keepdim`` is True, the dimensions of
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is False.
482 483 484
        ignore_nan: (bool, optional): Whether to ignore NaN of input Tensor.
            If ``ignore_nan`` is True, it will calculate nanquantile.
            Otherwise it will calculate quantile. Default is False.
485 486

    Returns:
487 488
        Tensor, results of quantile along ``axis`` of ``x``.
        In order to obtain higher precision, data type of results will be float64.
489
    """
490
    # Validate x
491 492
    if not isinstance(x, Variable):
        raise TypeError("input x should be a Tensor.")
493 494 495 496 497 498 499 500 501 502 503

    # Validate q
    if isinstance(q, (int, float)):
        q = [q]
    elif isinstance(q, (list, tuple)):
        if len(q) <= 0:
            raise ValueError("q should not be empty")
    else:
        raise TypeError("Type of q should be int, float, list or tuple.")

    # Validate axis
504
    dims = len(x.shape)
505
    out_shape = list(x.shape)
506 507 508 509 510 511 512 513 514
    if axis is None:
        x = paddle.flatten(x)
        axis = 0
        out_shape = [1] * dims
    else:
        if isinstance(axis, list):
            axis_src, axis_dst = [], []
            for axis_single in axis:
                if not isinstance(axis_single, int) or not (
515 516
                    axis_single < dims and axis_single >= -dims
                ):
517 518 519 520 521 522 523
                    raise ValueError(
                        "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
                    )
                if axis_single < 0:
                    axis_single = axis_single + dims
                axis_src.append(axis_single)
                out_shape[axis_single] = 1
524

525 526
            axis_dst = list(range(-len(axis), 0))
            x = paddle.moveaxis(x, axis_src, axis_dst)
527 528 529 530 531 532
            if len(axis_dst) == 0:
                x = paddle.flatten(x)
                axis = 0
            else:
                x = paddle.flatten(x, axis_dst[0], axis_dst[-1])
                axis = axis_dst[0]
533 534 535 536 537 538 539 540
        else:
            if not isinstance(axis, int) or not (axis < dims and axis >= -dims):
                raise ValueError(
                    "Axis should be None, int, or a list, element should in range [-rank(x), rank(x))."
                )
            if axis < 0:
                axis += dims
            out_shape[axis] = 1
541 542

    mask = x.isnan()
543 544 545
    valid_counts = mask.logical_not().sum(
        axis=axis, keepdim=True, dtype='float64'
    )
546

547
    indices = []
548 549 550

    for q_num in q:
        if q_num < 0 or q_num > 1:
551
            raise ValueError("q should be in range [0, 1]")
552
        if in_dygraph_mode():
553 554 555 556
            q_num = paddle.to_tensor(q_num, dtype='float64')
        if ignore_nan:
            indices.append(q_num * (valid_counts - 1))
        else:
557
            # TODO: Use paddle.index_fill instead of where
558 559 560 561 562 563
            index = q_num * (valid_counts - 1)
            last_index = x.shape[axis] - 1
            nums = paddle.full_like(index, fill_value=last_index)
            index = paddle.where(mask.any(axis=axis, keepdim=True), nums, index)
            indices.append(index)

564 565
    sorted_tensor = paddle.sort(x, axis)

566
    outputs = []
567

568
    # TODO(chenjianye): replace the for-loop to directly take elements.
569 570 571
    for index in indices:
        indices_below = paddle.floor(index).astype(paddle.int32)
        indices_upper = paddle.ceil(index).astype(paddle.int32)
572 573 574 575 576 577 578 579 580 581 582 583
        tensor_upper = paddle.take_along_axis(
            sorted_tensor, indices_upper, axis=axis
        )
        tensor_below = paddle.take_along_axis(
            sorted_tensor, indices_below, axis=axis
        )
        weights = index - indices_below.astype('float64')
        out = paddle.lerp(
            tensor_below.astype('float64'),
            tensor_upper.astype('float64'),
            weights,
        )
584 585 586 587 588
        if not keepdim:
            out = paddle.squeeze(out, axis=axis)
        else:
            out = out.reshape(out_shape)
        outputs.append(out)
589 590 591

    if len(q) > 1:
        outputs = paddle.stack(outputs, 0)
592
    else:
593 594 595 596 597 598 599 600 601 602 603
        outputs = outputs[0]

    return outputs


def quantile(x, q, axis=None, keepdim=False):
    """
    Compute the quantile of the input along the specified axis.
    If any values in a reduced row are NaN, then the quantiles for that reduction will be NaN.

    Args:
604
        x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628
        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
            each q will be calculated and the first dimension of output is same to the number of ``q`` .
        axis (int|list, optional): The axis along which to calculate quantile. ``axis`` should be int or list of int.
            ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
            If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
            If ``axis`` is a list, quantile is calculated over all elements of given axises.
            If ``axis`` is None, quantile is calculated over all elements of ``x``. Default is None.
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
            in the output Tensor. If ``keepdim`` is True, the dimensions of
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is False.
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor, results of quantile along ``axis`` of ``x``.
        In order to obtain higher precision, data type of results will be float64.

    Examples:
        .. code-block:: python

            import paddle

629 630 631 632 633 634 635
            y = paddle.arange(0, 8 ,dtype="float32").reshape([4, 2])
            # Tensor(shape=[4, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
            #        [[0., 1.],
            #         [2., 3.],
            #         [4., 5.],
            #         [6., 7.]])

636
            y1 = paddle.quantile(y, q=0.5, axis=[0, 1])
637 638
            # Tensor(shape=[], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        3.50000000)
639 640

            y2 = paddle.quantile(y, q=0.5, axis=1)
641 642
            # Tensor(shape=[4], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [0.50000000, 2.50000000, 4.50000000, 6.50000000])
643 644

            y3 = paddle.quantile(y, q=[0.3, 0.5], axis=0)
645 646 647
            # Tensor(shape=[2, 2], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [[1.80000000, 2.80000000],
            #         [3.        , 4.        ]])
648

649
            y[0,0] = float("nan")
650
            y4 = paddle.quantile(y, q=0.8, axis=1, keepdim=True)
651 652 653 654 655
            # Tensor(shape=[4, 1], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [[nan       ],
            #         [2.80000000],
            #         [4.80000000],
            #         [6.80000000]])
656 657 658 659 660 661 662 663 664 665 666

    """
    return _compute_quantile(x, q, axis=axis, keepdim=keepdim, ignore_nan=False)


def nanquantile(x, q, axis=None, keepdim=False):
    """
    Compute the quantile of the input as if NaN values in input did not exist.
    If all values in a reduced row are NaN, then the quantiles for that reduction will be NaN.

    Args:
667
        x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691
        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
            each q will be calculated and the first dimension of output is same to the number of ``q`` .
        axis (int|list, optional): The axis along which to calculate quantile. ``axis`` should be int or list of int.
            ``axis`` should be in range [-D, D), where D is the dimensions of ``x`` .
            If ``axis`` is less than 0, it works the same way as :math:`axis + D`.
            If ``axis`` is a list, quantile is calculated over all elements of given axises.
            If ``axis`` is None, quantile is calculated over all elements of ``x``. Default is None.
        keepdim (bool, optional): Whether to reserve the reduced dimension(s)
            in the output Tensor. If ``keepdim`` is True, the dimensions of
            the output Tensor is the same as ``x`` except in the reduced
            dimensions(it is of size 1 in this case). Otherwise, the shape of
            the output Tensor is squeezed in ``axis`` . Default is False.
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor, results of quantile along ``axis`` of ``x``.
        In order to obtain higher precision, data type of results will be float64.

    Examples:
        .. code-block:: python

            import paddle

692
            x = paddle.to_tensor(
693
                [[0, 1, 2, 3, 4],
694 695 696
                    [5, 6, 7, 8, 9]],
                dtype="float32")
            x[0,0] = float("nan")
697 698

            y1 = paddle.nanquantile(x, q=0.5, axis=[0, 1])
699 700
            # Tensor(shape=[], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        5.)
701 702

            y2 = paddle.nanquantile(x, q=0.5, axis=1)
703 704
            # Tensor(shape=[2], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [2.50000000, 7.        ])
705 706

            y3 = paddle.nanquantile(x, q=[0.3, 0.5], axis=0)
707 708 709
            # Tensor(shape=[2, 5], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [[5.        , 2.50000000, 3.50000000, 4.50000000, 5.50000000],
            #         [5.        , 3.50000000, 4.50000000, 5.50000000, 6.50000000]])
710 711

            y4 = paddle.nanquantile(x, q=0.8, axis=1, keepdim=True)
712 713 714
            # Tensor(shape=[2, 1], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [[3.40000000],
            #         [8.20000000]])
715

716
            nan = paddle.full(shape=[2, 3], fill_value=float("nan"))
717
            y5 = paddle.nanquantile(nan, q=0.8, axis=1, keepdim=True)
718 719 720
            # Tensor(shape=[2, 1], dtype=float64, place=Place(cpu), stop_gradient=True,
            #        [[nan],
            #         [nan]])
721 722 723

    """
    return _compute_quantile(x, q, axis=axis, keepdim=keepdim, ignore_nan=True)