Numpy compatible multinomial

python/mxnet/_numpy_op_doc.py

@@ -109,3 +109,33 @@ def _np_repeat(a, repeats, axis=None):
  the given axis.
  """
  pass
+
+
+def _npi_multinomial(a):
+ """Draw samples from a multinomial distribution.
+
+ The multinomial distribution is a multivariate generalisation of the binomial distribution.
+ Take an experiment with one of ``p`` possible outcomes. An example of such an experiment is throwing a dice,
+ where the outcome can be 1 through 6. Each sample drawn from the distribution represents n such experiments.
+ Its values, ``X_i = [X_0, X_1, ..., X_p]``, represent the number of times the outcome was ``i``.
+
+
+ Parameters
+ ----------
+ n : int
+ Number of experiments.
+ pvals : sequence of floats, length p
+ Probabilities of each of the p different outcomes. These should sum to 1
+ (however, the last element is always assumed to account for the remaining
+ probability, as long as ``sum(pvals[:-1]) <= 1)``.
+ size : int or tuple of ints, optional
+ Output shape. If the given shape is, e.g., ``(m, n, k)``, then ``m * n * k`` sam-
+ ples are drawn. Default is None, in which case a single value is returned.
+
+ Returns
+ -------
+ out : ndarray
+ The drawn samples, of shape size, if that was provided. If not, the shape is ``(N,)``.
+ In other words, each entry ``out[i,j,...,:]`` is an N-dimensional value drawn from the distribution.
+ """
+ pass

python/mxnet/ndarray/numpy/random.py

@@ -17,11 +17,13 @@
 
 """Namespace for operators used in Gluon dispatched by F=ndarray."""
 from __future__ import absolute_import
+import numpy as np
 from ...base import numeric_types
 from ...context import current_context
+from ..ndarray import NDArray
 from . import _internal as _npi
 
-__all__ = ['uniform', 'normal']
+__all__ = ['uniform', 'normal', 'multinomial']
 
 
 def _random_helper(random, sampler, params, shape, dtype, ctx, out, kwargs):
@@ -135,3 +137,40 @@ def normal(loc=0.0, scale=1.0, size=None, **kwargs):
  out = kwargs.pop('out', None)
  return _random_helper(_npi.random_normal, None,
  [loc, scale], size, dtype, ctx, out, kwargs)
+
+
+def multinomial(n, pvals, size=None):
+ """Draw samples from a multinomial distribution.
+
+ The multinomial distribution is a multivariate generalisation of the binomial distribution.
+ Take an experiment with one of ``p`` possible outcomes. An example of such an experiment is throwing a dice,
+ where the outcome can be 1 through 6. Each sample drawn from the distribution represents n such experiments.
+ Its values, ``X_i = [X_0, X_1, ..., X_p]``, represent the number of times the outcome was ``i``.
+
+
+ Parameters
+ ----------
+ n : int
+ Number of experiments.
+ pvals : sequence of floats, length p
+ Probabilities of each of the p different outcomes. These should sum to 1
+ (however, the last element is always assumed to account for the remaining
+ probability, as long as ``sum(pvals[:-1]) <= 1)``.
+ size : int or tuple of ints, optional
+ Output shape. If the given shape is, e.g., ``(m, n, k)``, then ``m * n * k`` sam-
+ ples are drawn. Default is None, in which case a single value is returned.
+
+ Returns
+ -------
+ out : ndarray
+ The drawn samples, of shape size, if that was provided. If not, the shape is ``(N,)``.
+ In other words, each entry ``out[i,j,...,:]`` is an N-dimensional value drawn from the distribution.
+ """
+ if isinstance(pvals, NDArray):
+ return _npi.multinomial(pvals, pvals=None, n=n, size=size)
+ else:
+ if isinstance(pvals, np.ndarray):
+ pvals = pvals.tolist()
+ if any(isinstance(i, list) for i in pvals):
+ raise ValueError('object too deep for desired array')
+ return _npi.multinomial(n=n, pvals=pvals, size=size)

python/mxnet/numpy/random.py

@@ -98,3 +98,33 @@ def normal(loc=0.0, scale=1.0, size=None, **kwargs):
  This function currently does not support ``loc`` and ``scale`` as ndarrays.
  """
  return _mx_nd_np.random.normal(loc, scale, size, **kwargs)
+
+
+def multinomial(n, pvals, size=None, **kwargs):
+ """Draw samples from a multinomial distribution.
+
+ The multinomial distribution is a multivariate generalisation of the binomial distribution.
+ Take an experiment with one of ``p`` possible outcomes. An example of such an experiment is throwing a dice,
+ where the outcome can be 1 through 6. Each sample drawn from the distribution represents n such experiments.
+ Its values, ``X_i = [X_0, X_1, ..., X_p]``, represent the number of times the outcome was ``i``.
+
+
+ Parameters
+ ----------
+ n : int
+ Number of experiments.
+ pvals : sequence of floats, length p
+ Probabilities of each of the p different outcomes. These should sum to 1
+ (however, the last element is always assumed to account for the remaining
+ probability, as long as ``sum(pvals[:-1]) <= 1)``.
+ size : int or tuple of ints, optional
+ Output shape. If the given shape is, e.g., ``(m, n, k)``, then ``m * n * k`` sam-
+ ples are drawn. Default is None, in which case a single value is returned.
+
+ Returns
+ -------
+ out : ndarray
+ The drawn samples, of shape size, if that was provided. If not, the shape is ``(N,)``.
+ In other words, each entry ``out[i,j,...,:]`` is an N-dimensional value drawn from the distribution.
+ """
+ return _mx_nd_np.random.multinomial(n, pvals, size, **kwargs)

src/operator/numpy/random/np_multinomial_op.cc

@@ -0,0 +1,61 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2019 by Contributors
+ * \file np_multinomial_op.h
+ * \brief Operator for numpy sampling from multinomial distributions
+ */
+#include "./np_multinomial_op.h"
+
+namespace mxnet {
+namespace op {
+
+DMLC_REGISTER_PARAMETER(NumpyMultinomialParam);
+
+NNVM_REGISTER_OP(_npi_multinomial)
+.describe(R"code(Draw samples from a multinomial distribution. "
+"The multinomial distribution is a multivariate generalisation of the binomial distribution. "
+"Take an experiment with one of p possible outcomes. "
+"An example of such an experiment is throwing a dice, where the outcome can be 1 through 6. "
+"Each sample drawn from the distribution represents n such experiments. "
+"Its values, X_i = [X_0, X_1, ..., X_p], represent the number of times the outcome was i.
+)code")
+.set_num_inputs(
+ [](const nnvm::NodeAttrs& attrs) {
+ const NumpyMultinomialParam& param = nnvm::get<NumpyMultinomialParam>(attrs.parsed);
+ return param.pvals.has_value() ? 0U : 1U;
+ }
+)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<NumpyMultinomialParam>)
+.set_attr<mxnet::FInferShape>("FInferShape", NumpyMultinomialOpShape)
+.set_attr<nnvm::FInferType>("FInferType", NumpyMultinomialOpType)
+.set_attr<FResourceRequest>("FResourceRequest",
+ [](const nnvm::NodeAttrs& attrs) {
+ return std::vector<ResourceRequest>{
+ ResourceRequest::kRandom, ResourceRequest::kTempSpace};
+ })
+.set_attr<FCompute>("FCompute<cpu>", NumpyMultinomialForward<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
+.add_argument("a", "NDArray-or-Symbol", "Source input")
+.add_arguments(NumpyMultinomialParam::__FIELDS__());
+
+} // namespace op
+} // namespace mxnet

src/operator/numpy/random/np_multinomial_op.cu

@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2019 by Contributors
+ * \file np_multinomial_op.cu
+ * \brief Operator for numpy sampling from multinomial distributions
+ */
+#include "./np_multinomial_op.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_npi_multinomial)
+.set_attr<FCompute>("FCompute<gpu>", NumpyMultinomialForward<gpu>);
+
+} // namespace op
+} // namespace mxnet