CuPy and Dirichlet gives me TypeError: unsupported operand type(s) for +=: 'int' and 'tuple'

I simply want to create a random matrix A whose vectors are drawn from the Dirichlet distribution. The function works fine with numpy:

import numpy as np
A = np.random.dirichlet(np.ones(n), n)

When I do the same thing with cupy

import cupy as cp
A = cp.random.dirichlet(cp.ones(n), n)

I get the error below:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-1-45a4f64a8b6e> in <module>
      6 n = 10000 #Size of the square matrix
      7 
----> 8 A = cp.random.dirichlet(cp.ones(n), n)
      9 
     10 print("--- %s seconds ---" % (time.time() - start_time))

~\anaconda3\envs\tensorflow\lib\site-packages\cupy\random\_distributions.py in dirichlet(alpha, size, dtype)
    112     """
    113     rs = _generator.get_random_state()
--> 114     return rs.dirichlet(alpha, size, dtype)
    115 
    116 

~\anaconda3\envs\tensorflow\lib\site-packages\cupy\random\_generator.py in dirichlet(self, alpha, size, dtype)
    144             size = alpha.shape
    145         else:
--> 146             size += alpha.shape
    147         y = cupy.empty(shape=size, dtype=dtype)
    148         _kernels.standard_gamma_kernel(alpha, self._rk_seed, y)

TypeError: unsupported operand type(s) for +=: 'int' and 'tuple'

When the input is a numpy array like this

import cupy as cp
import numpy as np

A = cp.random.dirichlet(np.ones(n), n)

then I get the same error.

The alpha.shape from line 146 is (n,) when I check manually. Is it a cupy bug or am I missing something?

I'm using cupy-cuda101 version 8.5.0 for CUDA 10.1. Everything else that has to do with cupy and tensorflow works perfectly on my GPU (2080ti).

Solution

This is a bug in cupy which you should report on their GitHub.

They do not properly handle the case of an integer argument, despite the documentation. They require that you provide either a tuple or None. This is why you see the behavior you’re seeing. (If you provided a tuple (a, b), then the resulting shape would properly be (a, b, n).

The workaround here is to provide the shape you want as a length-1 tuple: (n,). Note that the comma is necessary.