@register_opt()
@local_optimizer([tensor.Rebroadcast])
def local_gpu_rebroadcast(node):
 if isinstance(node.op, tensor.Rebroadcast):
  (x,) = node.inputs
  if (x.owner and isinstance(x.owner.op, HostFromGpu)):
   gpu_x = x.owner.inputs[0]
   return [host_from_gpu(node.op(gpu_x))]

@ register_opt ( ) @ local_optimizer ( [ tensor . Rebroadcast ] ) def local_gpu_rebroadcast ( node ) : if isinstance ( node . op , tensor . Rebroadcast ) : ( x , ) = node . inputs if ( x . owner and isinstance ( x . owner . op , HostFromGpu ) ) : gpu_x = x . owner . inputs [ 0 ] return [ host_from_gpu ( node . op ( gpu_x ) ) ]
rebroadcast(host_from_gpu(x)) -> host_from_gpu(rebroadcast(x)) .