Repository: scipy/scipy

Issue:
improper handling of multidimensional input in stats.rv_sample
<!-- 
Thank you for taking the time to report a SciPy issue.
Please describe the issue in detail, and for bug reports
fill in the fields below. You can delete the sections that 
don't apply to your issue. You can view the final output
by clicking the preview button above.
-->

My issue is about ...

### Reproducing code example:
<!-- 
If you place your code between the triple backticks below, 
it will be marked as a code block automatically 
-->


```
from scipy import stats
import numpy as np

xk = np.arange(6).reshape((3, 2))
pk = (1/3, 1/3, 1/3)
custm = stats.rv_discrete(values=(xk, pk))
```
The implementation of `rv_sample` only checks the following:

```
if len(xk) != len(pk):
            raise ValueError("xk and pk need to have the same length.")
```

Afterwards, xk and pk are ravelled, leading to different lengths.

### Error message:

`IndexError: index 3 is out of bounds for size 3`

```
Traceback (most recent call last):

  File "<ipython-input-15-0a2191a05404>", line 3, in <module>
    custm = stats.rv_discrete(values=(xk, pk))

  File "c:\users\christoph\anaconda3\lib\site-packages\scipy\stats\_distn_infrastructure.py", line 3332, in __init__
    self.pk = np.take(np.ravel(pk), indx, 0)

  File "c:\users\christoph\anaconda3\lib\site-packages\numpy\core\fromnumeric.py", line 181, in take
    return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode)

  File "c:\users\christoph\anaconda3\lib\site-packages\numpy\core\fromnumeric.py", line 51, in _wrapfunc
    return getattr(obj, method)(*args, **kwds)
IndexError: index 3 is out of bounds for size 3
```

### Scipy/Numpy/Python version information:

```
import sys, scipy, numpy; print(scipy.__version__, numpy.__version__, sys.version_info)
1.1.0 1.15.2 sys.version_info(major=3, minor=5, micro=4, releaselevel='final', serial=0)
```




Discussion:
 
 ... and once this is fixed, most likely something's going to break down the line, too. Am labelling this as an enhancement.

EDIT: scratch that, pk is raveled. So it's just a bug.

Resolving patch:
diff --git a/scipy/stats/_distn_infrastructure.py b/scipy/stats/_distn_infrastructure.py
--- a/scipy/stats/_distn_infrastructure.py
+++ b/scipy/stats/_distn_infrastructure.py
@@ -2579,6 +2579,7 @@ class rv_discrete(rv_generic):
     values : tuple of two array_like, optional
         ``(xk, pk)`` where ``xk`` are integers with non-zero
         probabilities ``pk``  with ``sum(pk) = 1``.
+        ``xk`` and ``pk`` must have the same shape.
     inc : integer, optional
         Increment for the support of the distribution.
         Default is 1. (other values have not been tested)
@@ -3360,8 +3361,8 @@ def __init__(self, a=0, b=inf, name=None, badvalue=None,
 
         xk, pk = values
 
-        if len(xk) != len(pk):
-            raise ValueError("xk and pk need to have the same length.")
+        if np.shape(xk) != np.shape(pk):
+            raise ValueError("xk and pk must have the same shape.")
         if not np.allclose(np.sum(pk), 1):
             raise ValueError("The sum of provided pk is not 1.")
 


Test patch:
diff --git a/scipy/stats/tests/test_distributions.py b/scipy/stats/tests/test_distributions.py
--- a/scipy/stats/tests/test_distributions.py
+++ b/scipy/stats/tests/test_distributions.py
@@ -1203,6 +1203,21 @@ def test_bad_input(self):
         pk = [1, 2, 3]
         assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk)))
 
+    def test_shape_rv_sample(self):
+        # tests added for gh-9565
+
+        # mismatch of 2d inputs
+        xk, pk = np.arange(4).reshape((2, 2)), np.ones((2, 3)) / 6
+        assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk)))
+
+        # same number of elements, but shapes not compatible
+        xk, pk = np.arange(6).reshape((3, 2)), np.ones((2, 3)) / 6
+        assert_raises(ValueError, stats.rv_discrete, **dict(values=(xk, pk)))
+
+        # same shapes => no error
+        xk, pk = np.arange(6).reshape((3, 2)), np.ones((3, 2)) / 6
+        assert_equal(stats.rv_discrete(values=(xk, pk)).pmf(0), 1/6)
+
 
 class TestSkewNorm(object):
     def setup_method(self):