sum_arrays.py 文件源码

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,