提交 8e614dee authored 作者: Pascal Lamblin's avatar Pascal Lamblin

Add optimization to execute rebroadcast on GPU

上级 10b7cb91
......@@ -280,6 +280,15 @@ def local_gpu_shape(node):
return [gpu_shape(gpu_x)]
return False
@register_opt()
@local_optimizer([])
def local_gpu_rebroadcast(node):
'''rebroadcast(host_from_gpu(x)) -> host_from_gpu(rebroadcast(x))'''
if isinstance(node.op, tensor.Rebroadcast):
x, = node.inputs
if (x.owner and x.owner.op == host_fom_gpu):
gpu_x = x.owner.inputs[0]
return [host_from_gpu(node.op(gpu_x))]
def cast(x, dtype):
stype = theano.scalar.Scalar(dtype)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论