Repository: mars-project/mars

Issue:
[BUG] mt.random.choice acts wrong
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Our implementation about `mt.random.choice` is wrong, for instance.

```
In [20]: t = mt.random.choice(mt.arange(20, chunk_size=10), size=15, replace=False, chunk_size=8)                                                      

In [21]: t.execute()                                                            
Out[21]: array([ 2, 16,  4,  3,  7,  1, 11, 10, 16,  8, 18,  4, 14, 19, 11])
```

Some duplicated values happens.
Support sampling without replacement 
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

In sampling without replacement, each sample unit of the population has only one chance to be selected in the sample.

**Describe the solution you'd like**

Distributed Reservoir Sampling method could be used.




Discussion:

Resolving patch:
diff --git a/mars/tensor/random/choice.py b/mars/tensor/random/choice.py
--- a/mars/tensor/random/choice.py
+++ b/mars/tensor/random/choice.py
@@ -19,33 +19,39 @@
 import numpy as np
 
 from ... import opcodes as OperandDef
-from ...serialize import AnyField, KeyField, BoolField
-from ..core import TENSOR_TYPE
+from ...config import options
+from ...serialize import ValueType, AnyField, KeyField, BoolField, TupleField
+from ...tiles import TilesError
+from ...utils import check_chunks_unknown_shape, ceildiv
+from ..operands import TensorOperandMixin
+from ..core import TENSOR_TYPE, CHUNK_TYPE, TensorOrder
 from ..datasource import arange, array
-from .core import TensorRandomOperandMixin, TensorSimpleRandomData
+from ..utils import decide_chunk_sizes, recursive_tile, normalize_chunk_sizes, gen_random_seeds
+from ..array_utils import as_same_device, device
+from .core import TensorRandomOperand, RandomState
 
 
-class TensorChoice(TensorSimpleRandomData, TensorRandomOperandMixin):
-    __slots__ = '_a', '_size', '_replace', '_p'
-    _into_one_chunk_fields_ = ['_a', '_p']
-    _input_fields_ = ['_a', '_p']
+class TensorChoice(TensorRandomOperand, TensorOperandMixin):
     _op_type_ = OperandDef.RAND_CHOICE
 
     _a = AnyField('a')
+    _size = TupleField('size', ValueType.int64)
     _replace = BoolField('replace')
     _p = KeyField('p')
-    _func_name = 'choice'
 
-    def __init__(self, a=None, p=None, state=None, size=None, replace=None,
-                 dtype=None, gpu=None, **kw):
-        dtype = np.dtype(dtype) if dtype is not None else dtype
-        super().__init__(_a=a, _p=p, _state=state, _size=size, _replace=replace,
-                         _dtype=dtype, _gpu=gpu, **kw)
+    def __init__(self, a=None, size=None, replace=None, p=None,
+                 state=None, seed=None, dtype=None, gpu=None, **kw):
+        super().__init__(_a=a, _size=size, _replace=replace, _p=p,
+                         _state=state, _seed=seed, _dtype=dtype, _gpu=gpu, **kw)
 
     @property
     def a(self):
         return self._a
 
+    @property
+    def size(self):
+        return self._size
+
     @property
     def replace(self):
         return self._replace
@@ -56,13 +62,197 @@ def p(self):
 
     def _set_inputs(self, inputs):
         super()._set_inputs(inputs)
-        if self._a is not None:
+        if isinstance(self._a, (TENSOR_TYPE, CHUNK_TYPE)):
             self._a = self._inputs[0]
-        if self._p is not None:
+        if isinstance(self._p, (TENSOR_TYPE, CHUNK_TYPE)):
             self._p = self._inputs[-1]
 
     def __call__(self, a, p, chunk_size=None):
-        return self.new_tensor([a, p], None, raw_chunk_size=chunk_size)
+        inputs = []
+        if isinstance(a, TENSOR_TYPE):
+            inputs.append(a)
+        if isinstance(p, TENSOR_TYPE):
+            inputs.append(p)
+        return self.new_tensor(inputs, shape=self._size,
+                               raw_chunk_size=chunk_size,
+                               order=TensorOrder.C_ORDER)
+
+    @classmethod
+    def _tile_one_chunk(cls, op, a, p):
+        out = op.outputs[0]
+        chunk_op = op.copy().reset_key()
+        chunk_op._seed = gen_random_seeds(1, op.state)[0]
+        chunk_inputs = []
+        if isinstance(a, TENSOR_TYPE):
+            chunk_op._a = a.chunks[0]
+            chunk_inputs.append(chunk_op.a)
+        else:
+            chunk_op._a = a
+        if isinstance(p, TENSOR_TYPE):
+            chunk_op._p = p.chunks[0]
+            chunk_inputs.append(chunk_op.p)
+        else:
+            chunk_op._p = p
+        chunk = chunk_op.new_chunk(chunk_inputs, shape=out.shape,
+                                   index=(0,) * out.ndim,
+                                   order=out.order)
+
+        new_op = op.copy()
+        return new_op.new_tensors(op.inputs, shape=out.shape,
+                                  order=out.order, chunks=[chunk],
+                                  nsplits=tuple((s,) for s in out.shape))
+
+    @classmethod
+    def _tile_sample_with_replacement(cls, op, a, nsplits):
+        out_shape = tuple(sum(ns) for ns in nsplits)
+        out_size = np.prod(out_shape).item()
+        most_chunk_size = np.prod([max(ns) for ns in nsplits]).item()
+
+        is_a_int = False
+        if isinstance(a, Integral):
+            is_a_int = True
+            a_size = a
+        else:
+            a = array(a)
+            a_size = a.size
+
+        rs = RandomState.from_numpy(op.state)
+
+        if is_a_int:
+            # the indices is just the result
+            ret = rs.randint(a_size, size=out_shape, chunk_size=nsplits)
+        else:
+            # gen indices first, need to be flattened
+            indices = rs.randint(a_size, size=out_size, chunk_size=most_chunk_size)
+            # get result via fancy indexing
+            ret = a[indices]
+            if len(out_shape) > 1:
+                # reshape back if out's ndim > 1
+                ret = ret.reshape(out_shape)
+            ret = ret.rechunk(nsplits)
+
+        return [recursive_tile(ret)]
+
+    @classmethod
+    def _tile_sample_without_replacement(cls, op, a, nsplits):
+        from ..base import searchsorted
+        from ..merge.stack import TensorStack
+        from ..indexing.getitem import TensorIndex
+
+        out = op.outputs[0]
+        out_shape = tuple(sum(ns) for ns in nsplits)
+        # to sample count
+        m = np.prod(out_shape).item()
+
+        if isinstance(a, Integral):
+            a_size = a
+            a = arange(a)
+        else:
+            a = array(a)
+            a_size = a.size
+        a = a._inplace_tile()
+
+        if any(cs < m for cs in a.nsplits[0]):
+            # make sure all chunk > m
+            n_chunk = min(max(a.size // (m + 1), 1), a.chunk_shape[0])
+            chunk_size = ceildiv(a.size, n_chunk)
+            chunk_sizes = normalize_chunk_sizes(a.size, chunk_size)[0]
+            if chunk_sizes[-1] < m and len(chunk_sizes) > 1:
+                # the last chunk may still less than m
+                # merge it into previous one
+                chunk_sizes[-2] += chunk_sizes[-1]
+                chunk_sizes = chunk_sizes[:-1]
+            a = a.rechunk({0: chunk_sizes})._inplace_tile()
+            if len(chunk_sizes) == 1:
+                return cls._tile_one_chunk(op, a, None)
+
+        # for each chunk in a, do regular sampling
+        sampled_chunks = []
+        sample_seeds = gen_random_seeds(len(a.chunks), op.state)
+        for seed, chunk in zip(sample_seeds, a.chunks):
+            chunk_op = op.copy().reset_key()
+            chunk_op._a = chunk
+            chunk_op._size = (m,)
+            chunk_op._seed = seed
+            sampled_chunk = chunk_op.new_chunk([chunk], shape=(m,),
+                                               order=out.order,
+                                               index=chunk.index)
+            sampled_chunks.append(sampled_chunk)
+
+        if len(sampled_chunks) == 1:
+            out_chunk = sampled_chunks[0]
+        else:
+            stacked_chunk = TensorStack(axis=0, dtype=sampled_chunks[0].dtype).new_chunk(
+                sampled_chunks, shape=(len(a.chunks), m), order=TensorOrder.C_ORDER)
+
+            # gen indices with length m from 0...a.size
+            state = RandomState.from_numpy(op.state)
+            indices = state.randint(a_size, size=(m,))
+            cum_offsets = np.cumsum(a.nsplits[0])
+            ind = recursive_tile(searchsorted(cum_offsets, indices, side='right'))
+            ind_chunk = ind.chunks[0]
+
+            # do fancy index to find result
+            indexes = (ind_chunk, arange(m)._inplace_tile().chunks[0])
+            out_chunk = TensorIndex(dtype=stacked_chunk.dtype, indexes=indexes).new_chunk(
+                [stacked_chunk] + list(indexes), shape=(m,), order=TensorOrder.C_ORDER)
+
+        ret = op.copy().new_tensor(op.inputs, shape=(m,), order=out.order,
+                                   nsplits=((m,),), chunks=[out_chunk])
+        if len(out_shape) > 0:
+            ret = ret.reshape(out_shape)._inplace_tile()
+        ret = ret.rechunk(nsplits)._inplace_tile()
+        return [recursive_tile(ret)]
+
+    @classmethod
+    def tile(cls, op):
+        check_chunks_unknown_shape(op.inputs, TilesError)
+
+        out = op.outputs[0]
+        chunk_size = out.extra_params.raw_chunk_size or options.tensor.chunk_size
+        nsplits = decide_chunk_sizes(out.shape, chunk_size, out.dtype.itemsize)
+        inputs = op.inputs
+
+        a, p = op.a, op.p
+        if p is not None:
+            # we cannot handle p in a parallel fashion
+            inputs = []
+            if isinstance(a, TENSOR_TYPE):
+                a = a.rechunk(a.shape)._inplace_tile()
+                inputs.append(a)
+            p = p.rechunk(p.shape)._inplace_tile()
+            inputs.append(p)
+
+            # ignore nsplits if p is specified
+            nsplits = ((s,) for s in out.shape)
+
+        # all inputs and outputs has 1 chunk
+        if all(len(inp.chunks) == 1 for inp in inputs) and \
+                all(len(ns) == 1 for ns in nsplits):
+            return cls._tile_one_chunk(op, a, p)
+
+        if op.replace:
+            return cls._tile_sample_with_replacement(op, a, nsplits)
+        else:
+            return cls._tile_sample_without_replacement(op, a, nsplits)
+
+    @classmethod
+    def execute(cls, ctx, op):
+        inputs, device_id, xp = as_same_device(
+            [ctx[inp.key] for inp in op.inputs], device=op.device, ret_extra=True)
+        if isinstance(op.a, CHUNK_TYPE):
+            a = inputs[0]
+        else:
+            a = op.a
+        if isinstance(op.p, CHUNK_TYPE):
+            p = inputs[-1]
+        else:
+            p = op.p
+
+        with device(device_id):
+            rs = xp.random.RandomState(op.seed)
+            ctx[op.outputs[0].key] = rs.choice(
+                a, size=op.size, replace=op.replace, p=p)
 
 
 def choice(random_state, a, size=None, replace=True, p=None, chunk_size=None, gpu=None):
@@ -147,14 +337,13 @@ def choice(random_state, a, size=None, replace=True, p=None, chunk_size=None, gp
     if isinstance(a, Integral):
         if a <= 0:
             raise ValueError('a must be greater than 0')
-        a = arange(a, chunk_size=chunk_size)
+        a_size = a
         dtype = np.random.choice(1, size=(), p=np.array([1]) if p is not None else p).dtype
     else:
-        if not isinstance(a, TENSOR_TYPE):
-            a = np.asarray(a)
-        a = array(a, chunk_size=a.size).rechunk(a.size)  # do rechunk if a is already a tensor
+        a = array(a)
         if a.ndim != 1:
             raise ValueError('a must be one dimensional')
+        a_size = a.size
         dtype = a.dtype
 
     if p is not None:
@@ -163,12 +352,11 @@ def choice(random_state, a, size=None, replace=True, p=None, chunk_size=None, gp
             if not np.isclose(p.sum(), 1, rtol=1e-7, atol=0):
                 raise ValueError('probabilities do not sum to 1')
             p = array(p, chunk_size=p.size)
-        else:
-            p = p.rechunk(p.size)
         if p.ndim != 1:
             raise ValueError('p must be one dimensional')
 
     if size is None:
+        size = ()
         length = 1
     else:
         try:
@@ -176,7 +364,7 @@ def choice(random_state, a, size=None, replace=True, p=None, chunk_size=None, gp
             length = np.prod(size)
         except TypeError:
             length = size
-    if replace is False and length > a.size:
+    if replace is False and length > a_size:
         raise ValueError("Cannot take a larger sample than population when 'replace=False'")
 
     size = random_state._handle_size(size)


Test patch:
diff --git a/mars/tensor/random/tests/test_random.py b/mars/tensor/random/tests/test_random.py
--- a/mars/tensor/random/tests/test_random.py
+++ b/mars/tensor/random/tests/test_random.py
@@ -79,6 +79,29 @@ def testChoice(self):
         t = choice(5, 3, replace=False)
         self.assertEqual(t.shape, (3,))
 
+        with self.assertRaises(ValueError):
+            choice(-1)
+
+        # a should be 1-d
+        with self.assertRaises(ValueError):
+            choice(np.random.rand(2, 2))
+
+        # p sum != 1
+        with self.assertRaises(ValueError):
+            choice(np.random.rand(3), p=[0.2, 0.2, 0.2])
+
+        # p should b 1-d
+        with self.assertRaises(ValueError):
+            choice(np.random.rand(3), p=[[0.2, 0.6, 0.2]])
+
+        # replace=False, choice size cannot be greater than a.size
+        with self.assertRaises(ValueError):
+            choice(np.random.rand(10), 11, replace=False)
+
+        # replace=False, choice size cannot be greater than a.size
+        with self.assertRaises(ValueError):
+            choice(np.random.rand(10), (3, 4), replace=False)
+
     def testMultivariateNormal(self):
         mean = [0, 0]
         cov = [[1, 0], [0, 100]]
diff --git a/mars/tensor/random/tests/test_random_execute.py b/mars/tensor/random/tests/test_random_execute.py
--- a/mars/tensor/random/tests/test_random_execute.py
+++ b/mars/tensor/random/tests/test_random_execute.py
@@ -18,9 +18,10 @@
 
 import numpy as np
 
+from mars import tensor
 from mars.tensor.datasource import tensor as from_ndarray
+from mars.tiles import get_tiled
 from mars.lib.sparse.core import issparse
-from mars import tensor
 from mars.utils import ignore_warning
 from mars.tests.core import ExecutorForTest
 
@@ -117,51 +118,79 @@ def testSampleExecution(self):
             self.assertTrue(np.array_equal(res, np.random.RandomState(0).random_sample(size=(5, 5))))
 
     def testChoiceExecution(self):
-        arr = tensor.random.choice(5, size=3, chunk_size=1)
-        self.assertEqual(self.executor.execute_tensor(arr, concat=True)[0].shape, (3,))
-
-        arr = tensor.random.choice(5, size=(15,), chunk_size=5).tiles()
-        for chunk in arr.chunks:
-            chunk.op._seed = 0
-
-        for res in self.executor.execute_tensor(arr):
-            self.assertTrue(np.array_equal(res, np.random.RandomState(0).choice(5, size=(5,))))
-
-        arr = tensor.random.choice([1, 4, 9], size=3, chunk_size=1)
-        self.assertEqual(self.executor.execute_tensor(arr, concat=True)[0].shape, (3,))
-
-        arr = tensor.random.choice([1, 4, 9], size=(15,), chunk_size=5).tiles()
-        for chunk in arr.chunks:
-            chunk.op._seed = 0
-
-        for res in self.executor.execute_tensor(arr):
-            self.assertTrue(np.array_equal(res, np.random.RandomState(0).choice([1, 4, 9], size=(5,))))
-
-        with self.assertRaises(ValueError):
-            tensor.random.choice([1, 3, 4], size=5, replace=False, chunk_size=2)
-
-        arr = tensor.random.choice([1, 4, 9], size=3, replace=False, chunk_size=1)
-        self.assertEqual(self.executor.execute_tensor(arr, concat=True)[0].shape, (3,))
-
-        arr = tensor.random.choice([1, 4, 9], size=(3,), replace=False, chunk_size=1).tiles()
-        for chunk in arr.chunks:
-            chunk.op._seed = 0
-
-        for res in self.executor.execute_tensor(arr):
-            self.assertTrue(
-                np.array_equal(res, np.random.RandomState(0).choice([1, 4, 9], size=(1,), replace=False)))
-
-        arr = tensor.random.choice([1, 4, 9], size=3, p=[.2, .5, .3], chunk_size=1)
-        self.assertEqual(self.executor.execute_tensor(arr, concat=True)[0].shape, (3,))
-
-        arr = tensor.random.choice([1, 4, 9], size=(15,), chunk_size=5, p=[.2, .5, .3]).tiles()
-        for chunk in arr.chunks:
-            chunk.op._seed = 0
-
-        for res in self.executor.execute_tensor(arr):
-            self.assertTrue(
-                np.array_equal(res, np.random.RandomState(0).choice([1, 4, 9], size=(5,),
-                                                                    p=[.2, .5, .3])))
+        # test 1 chunk, get integer
+        a = tensor.random.RandomState(0).choice(10)
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        self.assertTrue(np.all((res >= 0) & (res < 10)))
+        seed = get_tiled(a).chunks[0].op.seed
+        expected = np.random.RandomState(seed).choice(10)
+        np.testing.assert_array_equal(res, expected)
+
+        # test 1 chunk, integer
+        a = tensor.random.RandomState(0).choice(10, (4, 3))
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        self.assertTrue(np.all((res >= 0) & (res < 10)))
+        b = tensor.random.RandomState(0).choice(10, (4, 3))
+        res2 = self.executor.execute_tensor(b, concat=True)[0]
+        np.testing.assert_array_equal(res, res2)
+        a_tiled = get_tiled(a)
+        seed = a_tiled.chunks[0].op.seed
+        expected = np.random.RandomState(seed).choice(10, (4, 3))
+        np.testing.assert_array_equal(res, expected)
+
+        # test 1 chunk, ndarray
+        raw = np.random.RandomState(0).rand(10)
+        a = tensor.random.RandomState(0).choice(raw, (4, 3))
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        seed = get_tiled(a).chunks[0].op.seed
+        expected = np.random.RandomState(seed).choice(raw, (4, 3))
+        np.testing.assert_array_equal(res, expected)
+
+        # test with replacement, integer
+        a = tensor.random.RandomState(0).choice(20, (7, 4), chunk_size=4)
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        self.assertTrue(np.all((res >= 0) & (res < 20)))
+        b = tensor.random.RandomState(0).choice(20, (7, 4), chunk_size=4)
+        res2 = self.executor.execute_tensor(b, concat=True)[0]
+        np.testing.assert_array_equal(res, res2)
+
+        # test with replacement, ndarray
+        raw = np.random.RandomState(0).rand(20)
+        t = tensor.array(raw, chunk_size=8)
+        a = tensor.random.RandomState(0).choice(t, (7, 4), chunk_size=4)
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        self.assertTrue(np.all((res >= 0) & (res < 20)))
+        b = tensor.random.RandomState(0).choice(t, (7, 4), chunk_size=4)
+        res2 = self.executor.execute_tensor(b, concat=True)[0]
+        np.testing.assert_array_equal(res, res2)
+
+        # test without replacement, integer
+        a = tensor.random.RandomState(0).choice(100, (7, 2), chunk_size=2, replace=False)
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        self.assertTrue(np.all((res >= 0) & (res < 100)))
+        self.assertTrue(len(np.unique(res)), a.size)
+        b = tensor.random.RandomState(0).choice(100, (7, 2), chunk_size=2, replace=False)
+        res2 = self.executor.execute_tensor(b, concat=True)[0]
+        np.testing.assert_array_equal(res, res2)
+
+        # test without replacement, ndarray
+        raw = np.random.RandomState(0).rand(100)
+        t = tensor.array(raw, chunk_size=47)
+        a = tensor.random.RandomState(0).choice(t, (7, 2), chunk_size=2, replace=False)
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        self.assertTrue(np.all((res >= 0) & (res < 100)))
+        self.assertTrue(len(np.unique(res)), a.size)
+        b = tensor.random.RandomState(0).choice(t, (7, 2), chunk_size=2, replace=False)
+        res2 = self.executor.execute_tensor(b, concat=True)[0]
+        np.testing.assert_array_equal(res, res2)
+
+        # test p
+        raw = np.random.RandomState(0).rand(5)
+        p = [0.3, 0.2, 0.1, 0.3, 0.1]
+        a = tensor.random.RandomState(0).choice(raw, 3, p=p)
+        res = self.executor.execute_tensor(a, concat=True)[0]
+        expected = np.random.RandomState(get_tiled(a).chunks[0].op.seed).choice(raw, 3, p=p)
+        np.testing.assert_array_equal(res, expected)
 
     def testSparseRandintExecution(self):
         size_executor = ExecutorForTest(sync_provider_type=ExecutorForTest.SyncProviderType.MOCK)