def forward_gpu(self, inputs):
n = len(inputs)
ptrs = cuda.cupy.asarray([x.data.ptr for x in inputs],
dtype=cuda.cupy.int64)
y = cuda.elementwise(
'T x0, int64 xs, int32 n_xs',
'T y',
'float** xs_ = (float**) xs;'
'y = 0;'
'for (size_t j = 0; j < n_xs; ++j) {'
' y += xs_[j][i];'
'}',
'sum_arrays'.format(n))(inputs[0], ptrs.data.ptr, len(ptrs))
return y,
评论列表
文章目录