提交 607102c6 authored 作者: James Bergstra's avatar James Bergstra

made GpuAlloc handle failure of CudaNdarray copy

上级 027d5c62
...@@ -1838,6 +1838,8 @@ class GpuAlloc(Op): ...@@ -1838,6 +1838,8 @@ class GpuAlloc(Op):
#the optimizer will remove them. #the optimizer will remove them.
v = as_cuda_ndarray_variable(value) v = as_cuda_ndarray_variable(value)
sh = [tensor.as_tensor_variable(s) for s in shape] sh = [tensor.as_tensor_variable(s) for s in shape]
if v.ndim != len(shape):
raise TypeError('GpuAlloc requires value of same dimensions as shape', value, len(shape))
bcast = [] bcast = []
for s in sh: for s in sh:
...@@ -1860,6 +1862,7 @@ class GpuAlloc(Op): ...@@ -1860,6 +1862,7 @@ class GpuAlloc(Op):
out[0][...] = v # broadcast v to fill us up out[0][...] = v # broadcast v to fill us up
def c_code(self, node, name, inputs, (out,), sub): def c_code(self, node, name, inputs, (out,), sub):
fail = sub['fail']
value = inputs[0] value = inputs[0]
shps = inputs[1:] shps = inputs[1:]
nd = len(shps) nd = len(shps)
...@@ -1875,7 +1878,10 @@ class GpuAlloc(Op): ...@@ -1875,7 +1878,10 @@ class GpuAlloc(Op):
%(out)s= (CudaNdarray*)CudaNdarray_new_null(); %(out)s= (CudaNdarray*)CudaNdarray_new_null();
CudaNdarray_alloc_contiguous(%(out)s, %(nd)s, dims); CudaNdarray_alloc_contiguous(%(out)s, %(nd)s, dims);
} }
CudaNdarray_CopyFromCudaNdarray(%(out)s, %(value)s, true); if (CudaNdarray_CopyFromCudaNdarray(%(out)s, %(value)s, true))
{
%(fail)s;
}
"""%locals() """%locals()
return str return str
...@@ -1886,7 +1892,7 @@ class GpuAlloc(Op): ...@@ -1886,7 +1892,7 @@ class GpuAlloc(Op):
return [None for i in inputs] return [None for i in inputs]
def c_code_cache_version(self): def c_code_cache_version(self):
return (1,) return (2,)
gpu_alloc = GpuAlloc() gpu_alloc = GpuAlloc()
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论