Skip to content

Commit 67b13de

Browse files
committed
fix comments
1 parent a3a2b86 commit 67b13de

File tree

6 files changed

+43
-16
lines changed

6 files changed

+43
-16
lines changed

python/tvm/relay/op/reduce.py

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ def argmax(data, axis=None, keepdims=False, exclude=False):
3939
4040
exclude : bool
4141
If `exclude` is true, reduction will be performed on the axes that are
42-
NOT in axis instead.
42+
NOT in axis instead.
4343
4444
Returns
4545
-------
@@ -69,7 +69,7 @@ def argmin(data, axis=None, keepdims=False, exclude=False):
6969
7070
exclude : bool
7171
If `exclude` is true, reduction will be performed on the axes that are
72-
NOT in axis instead.
72+
NOT in axis instead.
7373
7474
Returns
7575
-------
@@ -100,7 +100,7 @@ def sum(data, axis=None, keepdims=False, exclude=False):
100100
101101
exclude : bool
102102
If `exclude` is true, reduction will be performed on the axes that are
103-
NOT in axis instead.
103+
NOT in axis instead.
104104
105105
Returns
106106
-------
@@ -112,12 +112,12 @@ def sum(data, axis=None, keepdims=False, exclude=False):
112112

113113

114114
def all(data, axis=None, keepdims=False, exclude=False):
115-
"""Computes the logical and of array elements over given axes.
115+
"""Computes the logical AND of boolean array elements over given axes.
116116
117117
Parameters
118118
----------
119119
data : relay.Expr
120-
The input data
120+
The input boolean tensor
121121
122122
axis : None or int or tuple of int
123123
Axis or axes along which a sum is performed. The default, axis=None,
@@ -131,12 +131,33 @@ def all(data, axis=None, keepdims=False, exclude=False):
131131
132132
exclude : bool
133133
If `exclude` is true, reduction will be performed on the axes that are
134-
NOT in axis instead.
134+
NOT in axis instead.
135135
136136
Returns
137137
-------
138138
result : relay.Expr
139139
The computed result.
140+
141+
Examples
142+
--------
143+
.. code-block:: python
144+
145+
data = relay.Constant(tvm.nd.array([[[ True, True, True],
146+
[ True, True, True],
147+
[False, True, False]],
148+
[[ True, False, False],
149+
[ True, True, False],
150+
[False, True, True]]]))
151+
152+
relay.all(data, axis=1)
153+
# [[False, True, False],
154+
# [False, False, False]]
155+
156+
relay.all(data, axis=0)
157+
# [[ True, False, False],
158+
# [ True, True, False],
159+
# [False, True, False]]
160+
140161
"""
141162
axis = [axis] if axis and isinstance(axis, int) else axis
142163
return _make.all(data, axis, keepdims, exclude)
@@ -162,7 +183,7 @@ def max(data, axis=None, keepdims=False, exclude=False):
162183
163184
exclude : bool
164185
If `exclude` is true, reduction will be performed on the axes that are
165-
NOT in axis instead.
186+
NOT in axis instead.
166187
167188
Returns
168189
-------
@@ -194,7 +215,7 @@ def min(data, axis=None, keepdims=False, exclude=False):
194215
195216
exclude : bool
196217
If `exclude` is true, reduction will be performed on the axes that are
197-
NOT in axis instead.
218+
NOT in axis instead.
198219
199220
Returns
200221
-------
@@ -225,7 +246,7 @@ def mean(data, axis=None, keepdims=False, exclude=False):
225246
226247
exclude : bool
227248
If `exclude` is true, reduction will be performed on the axes that are
228-
NOT in axis instead.
249+
NOT in axis instead.
229250
230251
Returns
231252
-------
@@ -256,7 +277,7 @@ def prod(data, axis=None, keepdims=False, exclude=False):
256277
257278
exclude : bool
258279
If `exclude` is true, reduction will be performed on the axes that are
259-
NOT in axis instead.
280+
NOT in axis instead.
260281
261282
Returns
262283
-------

src/relay/op/tensor/reduce.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -364,7 +364,7 @@ Array<Tensor> AllCompute(const Attrs& attrs,
364364

365365

366366
RELAY_REGISTER_REDUCE_OP("all")
367-
.describe(R"code(Computes the logical AND of array elements over given axes.
367+
.describe(R"code(Computes the logical AND of boolean array elements over given axes.
368368
369369
Example::
370370

tests/python/relay/test_op_level10.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,7 @@ def test_adaptive_pool2d():
249249
verify_adaptive_pool2d((1, 14, 56, 78), (34, 13), "max")
250250
verify_adaptive_pool2d((1, 5, 46, 97), (4, 96), "avg")
251251

252+
252253
if __name__ == "__main__":
253254
test_adaptive_pool2d()
254255
test_collapse_sum_like()

tests/python/relay/test_op_level4.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,7 @@ def test_where():
138138
def verify_reduce(funcs, data, axis, keepdims, exclude, output, dtype="float32"):
139139
test_func = funcs[0]
140140
ref_func = funcs[1]
141+
dtype = "bool" if ref_func in [np.all] else dtype
141142

142143
x = relay.var("x", relay.TensorType(data, dtype))
143144
z = test_func(x, axis, keepdims, exclude)
@@ -155,7 +156,9 @@ def verify_reduce(funcs, data, axis, keepdims, exclude, output, dtype="float32")
155156
return
156157

157158
func = relay.Function([x], z)
158-
x_data = np.random.uniform(size=data).astype(dtype)
159+
x_data = np.random.choice([True, False], size=data) if ref_func in [np.all] \
160+
else np.random.uniform(size=data).astype(dtype)
161+
159162
if ref_func in [np.sum]:
160163
ref_res = ref_func(x_data + 0, axis=axis, dtype=dtype, keepdims=keepdims)
161164
elif ref_func in [np.max, np.min, np.mean, np.prod]:
@@ -194,6 +197,7 @@ def _wrapper(data, axis=None, keepdims=False):
194197
[relay.min, np.min],
195198
[relay.mean, np.mean],
196199
[relay.prod, np.prod],
200+
[relay.all, np.all],
197201
[relay.argmin, _with_keepdims(np.argmin)],
198202
[relay.argmax, _with_keepdims(np.argmax)]]:
199203
verify_reduce(func, (d1, d2, d3, d4), None, False, False, ())
@@ -203,6 +207,7 @@ def _wrapper(data, axis=None, keepdims=False):
203207
verify_reduce(func, (d1, d2, d3), (0, 1), True, False, (1, 1, d3))
204208
verify_reduce(func, (2, 3, 4), 1, True, False, (2, 1, 4))
205209
verify_reduce(func, (2, 3, 4), (1,), True, False, (2, 1, 4))
210+
verify_reduce(func, (2, 3, 4), -1, True, False, (2, 3, 1))
206211
verify_reduce(func, (2, 3, 4), (0, 1, 2), False, False, ())
207212
verify_reduce(func, (4, 4, 3), None, False, False, ())
208213
verify_reduce(func, (4, 4, 3), (0, 2), False, False, (4,))

topi/include/topi/reduction.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -372,9 +372,9 @@ inline Tensor collapse_sum(const Tensor& data, Array<Expr> target_shape) {
372372
* \brief Creates an operation that computes the logical AND of elements
373373
* over a given axis
374374
*
375-
* \param data The input tensor
376-
* \param axis The axis to perform logical AND over. If axis is empty, the
377-
* operation will perform logical AND over all elements of the array.
375+
* \param data The input boolean tensor
376+
* \param axis The axes to reduce. If axis is empty, the operation will
377+
* perform logical AND over all elements of the array.
378378
* \param keepdims If this is set to true, the axes which are reduced are
379379
* left in the result as dimensions with size one. This enables the result
380380
* to broadcast correctly against the input array.

topi/python/topi/reduction.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ def all(data, axis=None, keepdims=False):
7171
Parameters
7272
----------
7373
data : tvm.Tensor
74-
The input tvm tensor
74+
The input tvm boolean tensor
7575
7676
axis : None or int or tuple of int
7777
Axis or axes along which a logical AND is performed.

0 commit comments

Comments
 (0)