-
Notifications
You must be signed in to change notification settings - Fork 1.4k
/
ranges.py
151 lines (115 loc) · 4.32 KB
/
ranges.py
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import numpy
import cupy
from cupy import core
def arange(start, stop=None, step=1, dtype=None):
"""Returns an array with evenly spaced values within a given interval.
Values are generated within the half-open interval [start, stop). The first
three arguments are mapped like the ``range`` built-in function, i.e. start
and step are optional.
Args:
start: Start of the interval.
stop: End of the interval.
step: Step width between each pair of consecutive values.
dtype: Data type specifier. It is inferred from other arguments by
default.
Returns:
cupy.ndarray: The 1-D array of range values.
.. seealso:: :func:`numpy.arange`
"""
if dtype is None:
if any(numpy.dtype(type(val)).kind == 'f'
for val in (start, stop, step)):
dtype = float
else:
dtype = int
if stop is None:
stop = start
start = 0
size = int(numpy.ceil((stop - start) / step))
if size <= 0:
return cupy.empty((0,), dtype=dtype)
ret = cupy.empty((size,), dtype=dtype)
typ = numpy.dtype(dtype).type
_arange_ufunc(typ(start), typ(step), ret, dtype=dtype)
return ret
def linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None):
"""Returns an array with evenly-spaced values within a given interval.
Instead of specifying the step width like :func:`cupy.arange`, this
function requires the total number of elements specified.
Args:
start: Start of the interval.
stop: End of the interval.
num: Number of elements.
endpoint (bool): If ``True``, the stop value is included as the last
element. Otherwise, the stop value is omitted.
retstep (bool): If ``True``, this function returns (array, step).
Otherwise, it returns only the array.
dtype: Data type specifier. It is inferred from the start and stop
arguments by default.
Returns:
cupy.ndarray: The 1-D array of ranged values.
"""
if num < 0:
raise ValueError('linspace with num<0 is not supported')
if dtype is None:
# In actual implementation, only float is used
dtype = float
ret = cupy.empty((num,), dtype=dtype)
if num == 0:
step = float('nan')
elif num == 1:
ret.fill(start)
step = float('nan')
else:
div = (num - 1) if endpoint else num
step = float(stop - start) / div
stop = float(stop)
if step == 0.0:
# for underflow
_linspace_ufunc_underflow(start, stop - start, div, ret,
casting='unsafe')
else:
_linspace_ufunc(start, step, ret, casting='unsafe')
if endpoint:
ret[-1] = stop
if retstep:
return ret, step
else:
return ret
def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None):
"""Returns an array with evenly-spaced values on a log-scale.
Instead of specifying the step width like :func:`cupy.arange`, this
function requires the total number of elements specified.
Args:
start: Start of the interval.
stop: End of the interval.
num: Number of elements.
endpoint (bool): If ``True``, the stop value is included as the last
element. Otherwise, the stop value is omitted.
base (float): Base of the log space. The step sizes between the
elements on a log-scale are the same as ``base``.
dtype: Data type specifier. It is inferred from the start and stop
arguments by default.
Returns:
cupy.ndarray: The 1-D array of ranged values.
"""
y = linspace(start, stop, num=num, endpoint=endpoint)
if dtype is None:
return core.power(base, y)
return core.power(base, y).astype(dtype)
# TODO(okuta): Implement meshgrid
# mgrid
# ogrid
_arange_ufunc = core.create_ufunc(
'cupy_arange',
('bb->b', 'BB->B', 'hh->h', 'HH->H', 'ii->i', 'II->I', 'll->l', 'LL->L',
'qq->q', 'QQ->Q', 'ee->e', 'ff->f', 'dd->d'),
'out0 = in0 + i * in1')
_linspace_ufunc = core.create_ufunc(
'cupy_linspace',
('dd->d',),
'out0 = in0 + i * in1')
_linspace_ufunc_underflow = core.create_ufunc(
'cupy_linspace',
('ddd->d',),
'out0 = in0 + i * in1 / in2')