File size: 2,507 Bytes
008bd6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | from __future__ import annotations
import warnings
from cupy import _core
from cupy._core import fusion
from cupy._math import ufunc
from cupy._util import bf16_loop
def around(a, decimals=0, out=None):
"""Rounds to the given number of decimals.
Args:
a (cupy.ndarray): The source array.
decimals (int): Number of decimal places to round to (default: 0).
If decimals is negative, it specifies the number of positions to
the left of the decimal point.
out (cupy.ndarray): Output array.
Returns:
cupy.ndarray: Rounded array.
.. seealso:: :func:`numpy.around`
"""
if fusion._is_fusing():
return fusion._call_ufunc(
_core.core._round_ufunc, a, decimals, out=out)
a = _core.array(a, copy=False)
return a.round(decimals, out=out)
def round(a, decimals=0, out=None):
return around(a, decimals, out=out)
def round_(a, decimals=0, out=None):
warnings.warn('Please use `round` instead.', DeprecationWarning)
return around(a, decimals, out=out)
rint = ufunc.create_math_ufunc(
'rint', 1, 'cupy_rint',
'''Rounds each element of an array to the nearest integer.
.. seealso:: :data:`numpy.rint`
''')
def create_rounding_ufunc(name, op, doc):
return _core.create_ufunc(
name,
(
"?->?",
"b->b",
"B->B",
"h->h",
"H->H",
"i->i",
"I->I",
"l->l",
"L->L",
"q->q",
"Q->Q",
("e->e", op),
*bf16_loop(code=op),
("f->f", op),
("d->d", op),
),
"out0 = in0",
doc=doc,
)
floor = create_rounding_ufunc(
"cupy_floor",
"out0 = floor(in0)",
"""Rounds each element of an array to its floor integer.
.. seealso:: :data:`numpy.floor`
""",
)
ceil = create_rounding_ufunc(
"cupy_ceil",
"out0 = ceil(in0)",
"""Rounds each element of an array to its ceiling integer.
.. seealso:: :data:`numpy.ceil`
""",
)
trunc = create_rounding_ufunc(
"cupy_trunc",
"out0 = trunc(in0)",
"""Rounds each element of an array towards zero.
.. seealso:: :data:`numpy.trunc`
""",
)
fix = create_rounding_ufunc(
"cupy_fix",
"out0 = (in0 >= 0.0) ? floor(in0): ceil(in0)",
"""If given value x is positive, it return floor(x).
Else, it return ceil(x).
.. seealso:: :func:`numpy.fix`
""",
)
|