提交 dd544cba authored 作者: sentient07's avatar sentient07

Added docstrings and fixed two errors

上级 bb6829ba
...@@ -2,14 +2,13 @@ from __future__ import absolute_import, print_function, division ...@@ -2,14 +2,13 @@ from __future__ import absolute_import, print_function, division
import os import os
from theano import Apply, Op from theano import Apply, Op
from theano.tensor.extra_ops import CumsumOp from theano.tensor.extra_ops import CumsumOp
from .type import GpuArrayType
try: try:
from pygpu import gpuarray from pygpu import gpuarray
except ImportError: except ImportError:
pass pass
from .basic_ops import (as_gpuarray_variable, GpuKernelBase, Kernel, from .basic_ops import (as_gpuarray_variable, GpuKernelBase, Kernel)
infer_context_name)
from .opt import register_opt, op_lifter, register_opt2 from .opt import register_opt, op_lifter, register_opt2
...@@ -40,7 +39,6 @@ class GpuCumsum(GpuKernelBase, Op): ...@@ -40,7 +39,6 @@ class GpuCumsum(GpuKernelBase, Op):
def make_node(self, x): def make_node(self, x):
assert x.type.dtype == 'float32', "Only float32 supported for GpuCumSum" assert x.type.dtype == 'float32', "Only float32 supported for GpuCumSum"
x = as_gpuarray_variable(x, infer_context_name(x))
if x.ndim > GpuCumsum.SUPPORTED_NDIMS: if x.ndim > GpuCumsum.SUPPORTED_NDIMS:
raise NotImplementedError('Only cumsum on 1D, 2D and\ raise NotImplementedError('Only cumsum on 1D, 2D and\
...@@ -456,6 +454,8 @@ class GpuCumsum(GpuKernelBase, Op): ...@@ -456,6 +454,8 @@ class GpuCumsum(GpuKernelBase, Op):
@register_opt2([CumsumOp], 'fast_compile') @register_opt2([CumsumOp], 'fast_compile')
def use_gpu_cumsumop(op, ctx_name, inputs, outputs): def use_gpu_cumsumop(op, ctx_name, inputs, outputs):
if inputs[0].dtype == 'float32': if inputs[0].dtype == 'float32':
if isinstance(inputs[0].type, GpuArrayType):
return
axis = op.axis axis = op.axis
x = inputs[0] x = inputs[0]
......
...@@ -259,12 +259,19 @@ gpu_seqopt.register('InputToGpuArrayOptimizer', InputToGpuOptimizer(), ...@@ -259,12 +259,19 @@ gpu_seqopt.register('InputToGpuArrayOptimizer', InputToGpuOptimizer(),
class GraphToGPU(NavigatorOptimizer): class GraphToGPU(NavigatorOptimizer):
""" """
Transfer the graph as a whole to GPU instead of transfering node by node. Transfer the graph as a whole to GPU instead of transfering node by node.
Parameters
----------
local_optimizers_all : List or Set
The local optimizations to apply to a node.
local_optimizers_map : Dict
Dictionary object containing the mapping of Op to list of
LocalOptimizers.
""" """
def __init__(self, local_optimizers_all, local_optimizers_map): def __init__(self, local_optimizers_all, local_optimizers_map):
self.local_optimizers_all = local_optimizers_all self.local_optimizers_all = local_optimizers_all
self.local_optimizers_map = local_optimizers_map self.local_optimizers_map = local_optimizers_map
self.failure_callback = None
def add_requirements(self, fgraph): def add_requirements(self, fgraph):
fgraph.attach_feature(toolbox.ReplaceValidate()) fgraph.attach_feature(toolbox.ReplaceValidate())
...@@ -1082,7 +1089,7 @@ def local_gpua_careduce(op, context_name, inputs, outputs): ...@@ -1082,7 +1089,7 @@ def local_gpua_careduce(op, context_name, inputs, outputs):
for i in range(x.ndim): for i in range(x.ndim):
if i not in op.axis: if i not in op.axis:
out_shp.append(shape_i(x, i)) out_shp.append(shape_i(x, i))
unreshaped_reduce = GpuReshape(len(out_shp))( unreshaped_reduce = GpuReshape(len(out_shp))(reduce_reshaped_x,
tensor.stack(out_shp)) tensor.stack(out_shp))
else: else:
unreshaped_reduce = reduce_reshaped_x unreshaped_reduce = reduce_reshaped_x
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论