stat.py 29.4 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
18
from paddle import _C_ops, _ir_ops, ir
19
from paddle.framework import in_dynamic_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
            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

59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84
            >>> import paddle

            >>> 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.]]])
            >>> out1 = paddle.mean(x)
            >>> print(out1.numpy())
            12.5
            >>> out2 = paddle.mean(x, axis=-1)
            >>> print(out2.numpy())
            [[ 2.5  6.5 10.5]
             [14.5 18.5 22.5]]
            >>> out3 = paddle.mean(x, axis=-1, keepdim=True)
            >>> print(out3.numpy())
            [[[ 2.5]
              [ 6.5]
              [10.5]]
             [[14.5]
              [18.5]
              [22.5]]]
            >>> out4 = paddle.mean(x, axis=[0, 2])
            >>> print(out4.numpy())
            [ 8.5 12.5 16.5]
85
    """
86
    if in_dynamic_mode():
87
        return _C_ops.mean(x, axis, keepdim)
88
    else:
89 90
        if ir.core._use_new_ir_api():
            return _ir_ops.mean(x, axis, keepdim)
91 92 93 94 95 96
        reduce_all, axis = _get_reduce_axis_with_tensor(axis, x)
        check_variable_and_dtype(
            x,
            'x/input',
            ['uint16', 'float16', 'float32', 'float64'],
            'mean/reduce_mean',
97
        )
98 99 100 101 102 103 104 105 106 107 108
        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',
                )
109

110
        helper = LayerHelper('mean', **locals())
111

112 113 114 115 116 117 118 119 120
        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
121 122


123
def var(x, axis=None, unbiased=True, keepdim=False, name=None):
124
    """
125
    Computes the variance of ``x`` along ``axis`` .
126 127

    Args:
L
LoneRanger 已提交
128
        x (Tensor): The input Tensor with data type float16, float32, float64.
129 130 131 132
        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` .
133 134 135 136 137
            - 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`.
138 139

    Returns:
140
        Tensor, results of variance along ``axis`` of ``x``, with the same data type as ``x``.
141 142 143 144

    Examples:
        .. code-block:: python

145
            >>> import paddle
146

147 148 149 150 151 152 153
            >>> x = paddle.to_tensor([[1.0, 2.0, 3.0], [1.0, 4.0, 5.0]])
            >>> out1 = paddle.var(x)
            >>> print(out1.numpy())
            2.6666667
            >>> out2 = paddle.var(x, axis=1)
            >>> print(out2.numpy())
            [1.         4.3333335]
154
    """
155
    if not in_dynamic_mode():
L
LoneRanger 已提交
156 157 158
        check_variable_and_dtype(
            x, 'x', ['float16', 'float32', 'float64'], 'var'
        )
159 160

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

163
    dtype = x.dtype
164 165 166
    n = paddle.cast(paddle.numel(x), paddle.int64) / paddle.cast(
        paddle.numel(out), paddle.int64
    )
167
    n = n.astype(dtype)
168
    if unbiased:
169
        one_const = paddle.ones([], x.dtype)
170
        n = where(n > one_const, n - 1.0, one_const)
171
    n.stop_gradient = True
172 173 174
    out /= n
    return out

S
swtkiwi 已提交
175

176 177 178
def std(x, axis=None, unbiased=True, keepdim=False, name=None):
    """
    Computes the standard-deviation of ``x`` along ``axis`` .
L
Liufang Sang 已提交
179 180

    Args:
L
LoneRanger 已提交
181
        x (Tensor): The input Tensor with data type float16, float32, float64.
182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203
        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 已提交
204 205

    Returns:
206 207 208
        Tensor, results of standard-deviation along ``axis`` of ``x``, with the
        same data type as ``x``.

L
Liufang Sang 已提交
209 210 211
    Examples:
        .. code-block:: python

212
            >>> import paddle
213

214 215 216 217 218 219 220 221 222 223
            >>> x = paddle.to_tensor([[1.0, 2.0, 3.0], [1.0, 4.0, 5.0]])
            >>> out1 = paddle.std(x)
            >>> print(out1.numpy())
            1.6329932
            >>> out2 = paddle.std(x, unbiased=False)
            >>> print(out2.numpy())
            1.490712
            >>> out3 = paddle.std(x, axis=1)
            >>> print(out3.numpy())
            [1.       2.081666]
224

L
Liufang Sang 已提交
225
    """
226
    if not in_dynamic_mode():
L
LoneRanger 已提交
227 228 229
        check_variable_and_dtype(
            x, 'x', ['float16', 'float32', 'float64'], 'std'
        )
230 231
    out = var(**locals())
    return paddle.sqrt(out)
232 233 234 235


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

    Args:
239
        x (Tensor): The input Tensor, it's data type can be bool, float16, float32, float64, int32, int64, complex64, complex128.
240 241
        name (str, optional): Name for the operation (optional, default is None).
            For more information, please refer to :ref:`api_guide_Name`.
242 243

    Returns:
244
        Tensor: The number of elements for the input Tensor, whose shape is [].
245 246 247 248

    Examples:
        .. code-block:: python

249
            >>> import paddle
250

251 252 253 254
            >>> x = paddle.full(shape=[4, 5, 7], fill_value=0, dtype='int32')
            >>> numel = paddle.numel(x)
            >>> print(numel.numpy())
            140
255 256 257


    """
258
    if in_dynamic_mode():
259
        return _C_ops.numel(x)
260 261 262 263 264 265 266 267 268
    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 已提交
269 270


271
def nanmedian(x, axis=None, keepdim=False, name=None):
272 273 274 275 276 277 278
    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:
279
        x (Tensor): The input Tensor, it's data type can be int32, int64, float16, bfloat16, float32, float64.
280 281 282 283 284 285 286 287 288
        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
289
            the output Tensor is squeezed in ``axis`` . Default is False.
290 291 292 293 294 295 296 297 298
        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

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

302 303 304
            >>> y1 = x.nanmedian()
            >>> print(y1.numpy())
            2.0
305

306 307 308
            >>> y2 = x.nanmedian(0)
            >>> print(y2.numpy())
            [0.  1.5 2.5]
309

310 311 312
            >>> y3 = x.nanmedian(0, keepdim=True)
            >>> print(y3.numpy())
            [[0.  1.5 2.5]]
313

314 315 316
            >>> y4 = x.nanmedian((0, 1))
            >>> print(y4.numpy())
            2.0
317 318 319 320 321 322 323 324 325 326 327 328 329 330
    """
    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.")

    if axis is None:
        axis = []
    elif isinstance(axis, tuple):
        axis = list(axis)
    elif isinstance(axis, int):
        axis = [axis]

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

341 342 343 344 345 346 347 348 349 350 351
        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
352 353


Z
zhulei 已提交
354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
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

378
            >>> import paddle
Z
zhulei 已提交
379

380 381 382 383 384 385
            >>> x = paddle.arange(12).reshape([3, 4])
            >>> print(x)
            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 已提交
386

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

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

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

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

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

411 412 413
    if x.size == 0:
        raise ValueError("In median, the size of input x should not be 0.")

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

    if axis is None:
        is_flatten = True

Z
zhulei 已提交
427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442
    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(
443 444
            tensor_topk, axes=[axis], starts=[kth - 1], ends=[kth]
        ) + paddle.slice(tensor_topk, axes=[axis], starts=[kth], ends=[kth + 1])
Z
zhulei 已提交
445 446
        out_tensor = paddle.cast(out_tensor, dtype=dtype) / 2
    else:
447 448 449 450 451 452
        out_tensor = paddle.cast(
            paddle.slice(
                tensor_topk, axes=[axis], starts=[kth], ends=[kth + 1]
            ),
            dtype=dtype,
        )
453
    out_tensor = out_tensor + paddle.sum(
454 455
        paddle.cast(paddle.isnan(x), dtype=dtype) * x, axis=axis, keepdim=True
    )
456 457 458
    if is_flatten:
        if keepdim:
            out_tensor = out_tensor.reshape([1] * dims)
Z
zhulei 已提交
459
        else:
460
            out_tensor = out_tensor.reshape([])
Z
zhulei 已提交
461
    else:
462 463
        if not keepdim:
            out_tensor = out_tensor.squeeze(axis)
Z
zhulei 已提交
464
    return out_tensor
465 466


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

471
    Args:
472
        x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
473
        q (int|float|list): The q for calculate quantile, which should be in range [0, 1]. If q is a list,
474 475 476 477 478 479 480 481 482 483 484
            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.
485 486 487
        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.
488 489

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

    # 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
507
    dims = len(x.shape)
508
    out_shape = list(x.shape)
509 510 511 512 513 514 515 516 517
    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 (
518 519
                    axis_single < dims and axis_single >= -dims
                ):
520 521 522 523 524 525 526
                    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
527

528 529
            axis_dst = list(range(-len(axis), 0))
            x = paddle.moveaxis(x, axis_src, axis_dst)
530 531 532 533 534 535
            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]
536 537 538 539 540 541 542 543
        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
544 545

    mask = x.isnan()
546 547 548
    valid_counts = mask.logical_not().sum(
        axis=axis, keepdim=True, dtype='float64'
    )
549

550
    indices = []
551 552 553

    for q_num in q:
        if q_num < 0 or q_num > 1:
554
            raise ValueError("q should be in range [0, 1]")
555
        if in_dynamic_mode():
556 557 558 559
            q_num = paddle.to_tensor(q_num, dtype='float64')
        if ignore_nan:
            indices.append(q_num * (valid_counts - 1))
        else:
560
            # TODO: Use paddle.index_fill instead of where
561 562 563 564 565 566
            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)

567 568
    sorted_tensor = paddle.sort(x, axis)

569
    outputs = []
570

571
    # TODO(chenjianye): replace the for-loop to directly take elements.
572 573 574
    for index in indices:
        indices_below = paddle.floor(index).astype(paddle.int32)
        indices_upper = paddle.ceil(index).astype(paddle.int32)
575 576 577 578 579 580 581 582 583 584 585 586
        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,
        )
587 588 589 590 591
        if not keepdim:
            out = paddle.squeeze(out, axis=axis)
        else:
            out = out.reshape(out_shape)
        outputs.append(out)
592 593 594

    if len(q) > 1:
        outputs = paddle.stack(outputs, 0)
595
    else:
596 597 598 599 600 601 602 603 604 605 606
        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:
607
        x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629
        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

630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663
            >>> import paddle

            >>> y = paddle.arange(0, 8 ,dtype="float32").reshape([4, 2])
            >>> print(y)
            Tensor(shape=[4, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
            [[0., 1.],
             [2., 3.],
             [4., 5.],
             [6., 7.]])

            >>> y1 = paddle.quantile(y, q=0.5, axis=[0, 1])
            >>> print(y1)
            Tensor(shape=[], dtype=float64, place=Place(cpu), stop_gradient=True,
            3.50000000)

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

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

            >>> y[0,0] = float("nan")
            >>> y4 = paddle.quantile(y, q=0.8, axis=1, keepdim=True)
            >>> print(y4)
            Tensor(shape=[4, 1], dtype=float64, place=Place(cpu), stop_gradient=True,
            [[nan       ],
             [2.80000000],
             [4.80000000],
             [6.80000000]])
664 665 666 667 668 669 670 671 672 673 674

    """
    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:
675
        x (Tensor): The input Tensor, it's data type can be float32, float64, int32, int64.
676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697
        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

698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733
            >>> import paddle

            >>> x = paddle.to_tensor(
            ...     [[0, 1, 2, 3, 4],
            ...      [5, 6, 7, 8, 9]],
            ...     dtype="float32")
            >>> x[0,0] = float("nan")

            >>> y1 = paddle.nanquantile(x, q=0.5, axis=[0, 1])
            >>> print(y1)
            Tensor(shape=[], dtype=float64, place=Place(cpu), stop_gradient=True,
            5.)

            >>> y2 = paddle.nanquantile(x, q=0.5, axis=1)
            >>> print(y2)
            Tensor(shape=[2], dtype=float64, place=Place(cpu), stop_gradient=True,
            [2.50000000, 7.        ])

            >>> y3 = paddle.nanquantile(x, q=[0.3, 0.5], axis=0)
            >>> print(y3)
            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]])

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

            >>> nan = paddle.full(shape=[2, 3], fill_value=float("nan"))
            >>> y5 = paddle.nanquantile(nan, q=0.8, axis=1, keepdim=True)
            >>> print(y5)
            Tensor(shape=[2, 1], dtype=float64, place=Place(cpu), stop_gradient=True,
            [[nan],
             [nan]])
734 735 736

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