weighted_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)
        ws = cuda.cupy.asarray(self.weights, dtype=cuda.cupy.float32)
        y = cuda.elementwise(
            'T x0, int64 xs, raw W ws, int32 n_xs',
            'T y',
            'float** xs_ = (float**) xs;'
            'y = 0;'
            'for (size_t j = 0; j < n_xs; ++j) {'
            '  y += xs_[j][i] * ws[j];'
            '}',
            'weighted_sum_arrays'.format(n))(inputs[0],
                                             ptrs.data.ptr,
                                             ws,
                                             len(ptrs))
        return y,