An equivalent to :func:`numpy.random.choice`. Use this function if you: * Perform a non-uniform sampling (probability tensor is given). * Sample a small set from a very large population (ratio <5%) uniformly *without* replacement. * Have a backend tensor on hand and does not
(a, size, replace=True, prob=None)
| 19 | |
| 20 | |
| 21 | def choice(a, size, replace=True, prob=None): # pylint: disable=invalid-name |
| 22 | """An equivalent to :func:`numpy.random.choice`. |
| 23 | |
| 24 | Use this function if you: |
| 25 | |
| 26 | * Perform a non-uniform sampling (probability tensor is given). |
| 27 | * Sample a small set from a very large population (ratio <5%) uniformly |
| 28 | *without* replacement. |
| 29 | * Have a backend tensor on hand and does not want to convert it to numpy |
| 30 | back and forth. |
| 31 | |
| 32 | Compared to :func:`numpy.random.choice`, it is slower when replace is True |
| 33 | and is comparable when replace is False. It wins when the population is |
| 34 | very large and the number of draws are quite small (e.g., draw <5%). The |
| 35 | reasons are two folds: |
| 36 | |
| 37 | * When ``a`` is a large integer, it avoids creating a large range array as |
| 38 | numpy does. |
| 39 | * When draw ratio is small, it switches to a hashmap based implementation. |
| 40 | |
| 41 | It out-performs numpy for non-uniform sampling in general cases. |
| 42 | |
| 43 | Parameters |
| 44 | ---------- |
| 45 | a : 1-D tensor or int |
| 46 | If an ndarray, a random sample is generated from its elements. If an int, |
| 47 | the random sample is generated as if a were F.arange(a) |
| 48 | size : int or tuple of ints |
| 49 | Output shape. E.g., for size ``(m, n, k)``, then ``m * n * k`` samples are drawn. |
| 50 | replace : bool, optional |
| 51 | If true, sample with replacement. |
| 52 | prob : 1-D tensor, optional |
| 53 | The probabilities associated with each entry in a. |
| 54 | If not given the sample assumes a uniform distribution over all entries in a. |
| 55 | |
| 56 | Returns |
| 57 | ------- |
| 58 | samples : 1-D tensor |
| 59 | The generated random samples |
| 60 | """ |
| 61 | # TODO(minjie): support RNG as one of the arguments. |
| 62 | if isinstance(size, tuple): |
| 63 | num = np.prod(size) |
| 64 | else: |
| 65 | num = size |
| 66 | |
| 67 | if F.is_tensor(a): |
| 68 | population = F.shape(a)[0] |
| 69 | else: |
| 70 | population = a |
| 71 | |
| 72 | if prob is None: |
| 73 | prob = nd.NULL["int64"] |
| 74 | else: |
| 75 | prob = F.zerocopy_to_dgl_ndarray(prob) |
| 76 | |
| 77 | bits = 64 # index array is in 64-bit |
| 78 | chosen_idx = _CAPI_Choice( |
no test coverage detected