提交 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
import os
from theano import Apply, Op
from theano.tensor.extra_ops import CumsumOp
from .type import GpuArrayType
try:
from pygpu import gpuarray
except ImportError:
pass
from .basic_ops import (as_gpuarray_variable, GpuKernelBase, Kernel,
infer_context_name)
from .basic_ops import (as_gpuarray_variable, GpuKernelBase, Kernel)
from .opt import register_opt, op_lifter, register_opt2
......@@ -40,7 +39,6 @@ class GpuCumsum(GpuKernelBase, Op):
def make_node(self, x):
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:
raise NotImplementedError('Only cumsum on 1D, 2D and\
......@@ -456,6 +454,8 @@ class GpuCumsum(GpuKernelBase, Op):
@register_opt2([CumsumOp], 'fast_compile')
def use_gpu_cumsumop(op, ctx_name, inputs, outputs):
if inputs[0].dtype == 'float32':
if isinstance(inputs[0].type, GpuArrayType):
return
axis = op.axis
x = inputs[0]
......
......@@ -259,12 +259,19 @@ gpu_seqopt.register('InputToGpuArrayOptimizer', InputToGpuOptimizer(),
class GraphToGPU(NavigatorOptimizer):
"""
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):
self.local_optimizers_all = local_optimizers_all
self.local_optimizers_map = local_optimizers_map
self.failure_callback = None
def add_requirements(self, fgraph):
fgraph.attach_feature(toolbox.ReplaceValidate())
......@@ -1082,8 +1089,8 @@ def local_gpua_careduce(op, context_name, inputs, outputs):
for i in range(x.ndim):
if i not in op.axis:
out_shp.append(shape_i(x, i))
unreshaped_reduce = GpuReshape(len(out_shp))(
tensor.stack(out_shp))
unreshaped_reduce = GpuReshape(len(out_shp))(reduce_reshaped_x,
tensor.stack(out_shp))
else:
unreshaped_reduce = reduce_reshaped_x
return [unreshaped_reduce]
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论