提交 dccf2bc0 authored 作者: Adam Becker's avatar Adam Becker

mixed changes

- GpuTopKOp: uncomment c_code_cache_version - add idx_dtype to TopKOp docs
上级 ed797624
...@@ -50,10 +50,8 @@ class GpuTopKOp(GpuKernelBase, TopKOp): ...@@ -50,10 +50,8 @@ class GpuTopKOp(GpuKernelBase, TopKOp):
def c_header_dirs(self): def c_header_dirs(self):
return [os.path.dirname(__file__), pygpu.get_include()] return [os.path.dirname(__file__), pygpu.get_include()]
'''
def c_code_cache_version(self): def c_code_cache_version(self):
return (1,) return (1,)
'''
def gpu_kernels(self, node, nodename): def gpu_kernels(self, node, nodename):
# load kernel source # load kernel source
......
...@@ -313,6 +313,10 @@ class TopKOp(theano.Op): ...@@ -313,6 +313,10 @@ class TopKOp(theano.Op):
If ``True``, one output of the Op will return the indices on the given axis. If ``True``, one output of the Op will return the indices on the given axis.
idx_dtype: string
Specify output dtype, defaults to ``int64``, must be integer type.
Notes Notes
----- -----
- ``return_values`` and ``return_indices`` cannot be both ``False`` - ``return_values`` and ``return_indices`` cannot be both ``False``
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论