[Numpy] Port nd.random.multinomial to npx.sample_categorical

src/operator/random/sample_multinomial_op.cc

@@ -32,6 +32,7 @@ DMLC_REGISTER_PARAMETER(SampleMultinomialParam);
 
 NNVM_REGISTER_OP(_sample_multinomial)
 .add_alias("sample_multinomial")
+.add_alias("_npx__random_categorical")
 .describe(R"code(Concurrent sampling from multiple multinomial distributions.
 
 *data* is an *n* dimensional array whose last dimension has length *k*, where

tests/python/unittest/test_numpy_op.py

@@ -4543,6 +4543,33 @@ def hybrid_forward(self, F, mean, cov):
         assert list(desired_shape) == list(actual_shape)
 
 
+@with_seed()
+@use_np
+def test_npx_categorical():
+    class TestNumpyCategorical(HybridBlock):
+        def __init__(self, size=None):
+            super(TestNumpyCategorical, self).__init__()
+            self.size = size
+
+        def hybrid_forward(self, F, prob):
+            if self.size is None:
+                return F.npx.random.categorical(prob)
+            return F.npx.random.categorical(prob, shape=self.size)
+
+    batch_sizes = [(2,), (2, 3)]
+    event_shapes = [None, (10,), (10, 12)]
+    num_event = [2, 4, 10]
+    for batch_size, num_event, event_shape in itertools.product(batch_sizes, num_event, event_shapes):
+        for hybridize in [True, False]:
+            prob = np.ones(batch_size + (num_event,)) / num_event
+            net = TestNumpyCategorical(event_shape)
+            if hybridize:
+                net.hybridize()
+            mx_out = net(prob)
+            desired_shape = batch_size + event_shape if event_shape is not None else batch_size
+            assert mx_out.shape == desired_shape
+
+
 @with_seed()
 @use_np
 def test_random_seed():