Cast to int64
@@ -206,6 +206,7 @@ def multinomial(n, pvals, size=None):
p = len(pvals)
shape = size + (p,)
+ # atomicAdd only supports int32
ys = basic.zeros(shape, 'i')
if ys.size > 0:
xs = choice(p, p=pvals, size=n * m)
@@ -213,4 +214,4 @@ def multinomial(n, pvals, size=None):
'int64 x, int32 p, int32 n', 'raw int32 ys',
'atomicAdd(&ys[i / n * p + x], 1)',
'cupy_random_multinomial')(xs, p, n, ys)
- return ys
+ return ys.astype('l')
