提交 330dd345 authored 作者: Adam Becker's avatar Adam Becker

add ``sorted`` to TopKOp.__props__

上级 c5848291
......@@ -33,6 +33,7 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
def __init__(
self, axis=-1,
sorted=True,
idx_dtype='int64',
return_values=True,
return_indices=True
......@@ -40,6 +41,7 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
GpuKernelBase.__init__(self)
TopKOp.__init__(
self, axis=axis,
sorted=sorted,
idx_dtype=idx_dtype,
return_values=return_values,
return_indices=return_indices)
......@@ -205,12 +207,12 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
} else if (dims[%(axis)d] < odims[%(axis)d]){
PyErr_SetString(
PyExc_ValueError,
"topk: kth cannot larger than size on specified axis %(axis)d");
"topk: kth cannot be larger than the size of specified axis %(axis)d");
%(fail)s;
} else if (dims[%(axis)d] >= (1u << 31)) {
PyErr_SetString(
PyExc_ValueError,
"topk: on GPU, array size on specified axis cannot larger or equal than 2^31");
"topk: on GPU, array size of specified axis cannot larger or equal than 2^31");
%(fail)s;
}
%(prep_output)s
......@@ -250,18 +252,16 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
blk[0] = %(MAX_TPB)d / 2;
err = GpuKernel_call(
&k_topk_dense_large_%(nodename)s, 3,
grd, blk, 0,
args);
grd, blk, 0, args);
} else {
err = GpuKernel_call(
&k_topk_dense_%(nodename)s, 3,
grd, blk, 0,
args);
grd, blk, 0, args);
}
if (err != GA_NO_ERROR) {
PyErr_SetString(
PyExc_RuntimeError,
"gpu kernel topk_kernel failed to execute");
"topk: gpu kernel failed to execute");
%(fail)s;
}
}
......@@ -302,6 +302,7 @@ def local_gpua_topkop(op, ctx_name, inputs, outputs):
gpu_op = GpuTopKOp(
axis=axis,
sorted=op.sorted,
idx_dtype=op.idx_dtype,
return_values=rv,
return_indices=ri)
......
......@@ -297,6 +297,12 @@ class TopKOp(theano.Op):
idx_dtype: string
Specify output dtype for indices, defaults to ``int64``, must be integer type.
sorted: bool
NOTE: NOT IMPLEMENTED YET
Defaults to ``True``
If True, the result array would be sorted in descending order.
Notes
-----
......@@ -319,11 +325,6 @@ class TopKOp(theano.Op):
# TODO more params
'''
sorted: bool
Defaults to ``True``
If True, the result array would be sorted in descending order.
only_top_kth: bool
Defaults to ``False``
......@@ -335,11 +336,12 @@ class TopKOp(theano.Op):
# TODO add opt, if k==1, use max/min reduce
# also if k is axis size, just copy input tensor
# TODO add opt, to merge argtopk / topk
__props__ = ('axis', 'return_values', 'return_indices', 'idx_dtype')
__props__ = ('axis', 'sorted', 'return_values', 'return_indices', 'idx_dtype')
def __init__(
self,
axis=-1,
sorted=True,
idx_dtype='int64',
return_values=True,
return_indices=True
......@@ -354,16 +356,19 @@ class TopKOp(theano.Op):
'"idx_dtype" parameter must be an integer dtype, got "%s"' % idx_dtype)
if not (return_indices or return_values):
raise ValueError("Neither return_values nor return_indices is True, this isn't allowd")
raise ValueError("Neither return_values nor return_indices is True, this isn't allowed")
self.axis = axis
self.sorted=sorted
self.return_values = return_values
self.return_indices = return_indices
self.idx_dtype = idx_dtype
def __str__(self):
return '%(op)s{axis=%(axis)d}' % dict(
op=self.__class__.__name__, axis=self.axis)
return '%(op)s{axis=%(axis)d, sorted=%(sorted)s}' % dict(
op=self.__class__.__name__,
axis=self.axis,
sorted=self.sorted)
def make_node(self, inp, kth):
inp = theano.tensor.as_tensor_variable(inp)
......@@ -446,6 +451,7 @@ def topk(x, kth, axis=-1, sorted=True, idx_dtype='int64'):
If ``None``, works on flattened array.
sorted: bool
NOTE: NOT IMPLEMENTED YET
Defaults to ``True``
If True, the result array would be sorted in descending order.
......@@ -469,7 +475,10 @@ def topk(x, kth, axis=-1, sorted=True, idx_dtype='int64'):
if axis is None:
x = theano.tensor.flatten(x)
axis = 0
return TopKOp(axis=axis, idx_dtype=idx_dtype)(x, kth)[0]
return TopKOp(
axis=axis,
sorted=sorted,
idx_dtype=idx_dtype)(x, kth)[0]
def argtopk(x, kth, axis=-1, sorted=True, idx_dtype='int64'):
......@@ -517,6 +526,7 @@ def argtopk(x, kth, axis=-1, sorted=True, idx_dtype='int64'):
axis = 0
return TopKOp(
axis=axis,
sorted=sorted,
idx_dtype=idx_dtype)(x, kth)[1]
......@@ -539,4 +549,5 @@ def topk_and_argtopk(x, kth, axis=-1, sorted=True, idx_dtype='int64'):
axis = 0
return TopKOp(
axis=axis,
sorted=sorted,
idx_dtype=idx_dtype)(x, kth)
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论