提交 455705e7 authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

Add optimization to move mrg_uniform to the GPU.

上级 5d221f46
...@@ -1341,6 +1341,15 @@ class MRG_RandomStreams(object): ...@@ -1341,6 +1341,15 @@ class MRG_RandomStreams(object):
assert final_samples.dtype == dtype assert final_samples.dtype == dtype
return final_samples return final_samples
from theano.sandbox.gpuarray.opt import (register_opt as register_gpua,
op_lifter as gpua_lifter)
@register_gpua()
@gpua_lifter([mrg_uniform])
def local_gpua_mrg(node):
return GPUA_mrg_uniform.new(node.inputs[0], node.op.output_type.ndim,
node.op.output_type.dtype, node.inputs[1])
@local_optimizer([mrg_uniform]) @local_optimizer([mrg_uniform])
def mrg_random_make_inplace(node): def mrg_random_make_inplace(node):
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论