[Numpy] Add sampling method for bernoulli

python/mxnet/base.py

@@ -752,7 +752,7 @@ def write_all_str(module_file, module_all_list):
 _NP_OP_SUBMODULE_LIST = ['_random_', '_linalg_']
 
 _NP_EXT_OP_PREFIX = '_npx_'
-_NP_EXT_OP_SUBMODULE_LIST = ['_image_']
+_NP_EXT_OP_SUBMODULE_LIST = ['_image_', '_random_']
 
 _NP_INTERNAL_OP_PREFIX = '_npi_'
 

python/mxnet/ndarray/numpy_extension/__init__.py

@@ -19,6 +19,7 @@
 
 from . import _op
 from . import image
+from . import random
 from . import _register
 from ._op import *  # pylint: disable=wildcard-import
 

python/mxnet/ndarray/numpy_extension/random.py

@@ -0,0 +1,57 @@
+# 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.
+
+"""Namespace for operators used in Gluon dispatched by F=ndarray."""
+from __future__ import absolute_import
+from ...context import current_context
+from .. import _internal as _npi
+
+
+__all__ = ['bernoulli']
+
+
+def bernoulli(probs, logits, size, dtype, ctx, out):
+    """
+    Sampling from bernoulli distribution.
+    """
+    from ...numpy import ndarray as np_ndarray
+    tensor_type_name = np_ndarray
+    if (probs is None) == (logits is None):
+        raise ValueError(
+            "Either `probs` or `logits` must be specified, but not both.")
+    if dtype is None:
+        dtype = 'float32'
+    if ctx is None:
+        ctx = current_context()
+    if size == ():
+        size = None
+    if probs is not None:
+        is_tensor = isinstance(probs, tensor_type_name)
+        if is_tensor:
+            return _npi.bernoulli(probs, probs=None, logits=None, is_logit=False,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)
+        else:
+            return _npi.bernoulli(probs=probs, logits=None, is_logit=False,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)
+    else:
+        is_tensor = isinstance(logits, tensor_type_name)
+        if is_tensor:
+            return _npi.bernoulli(logits, probs=None, logits=None, is_logit=True,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)
+        else:
+            return _npi.bernoulli(probs=None, logits=logits, is_logit=True,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)

python/mxnet/numpy_extension/__init__.py

@@ -22,12 +22,12 @@
 from __future__ import absolute_import
 from . import _op
 from . import image
+from . import random  # pylint: disable=wildcard-import
 from . import _register
 from ._op import *  # pylint: disable=wildcard-import
 from ..context import *  # pylint: disable=wildcard-import
 from ..util import is_np_shape, is_np_array, set_np, reset_np, get_cuda_compute_capability
 from ..ndarray import waitall
 from .utils import *  # pylint: disable=wildcard-import
-from . import random  # pylint: disable=wildcard-import
 
 __all__ = []

python/mxnet/numpy_extension/random.py

@@ -19,9 +19,10 @@
 
 from __future__ import absolute_import
 from .. import random as _mx_rand
+from ..ndarray import numpy_extension as _mx_nd_npx
 
 
-__all__ = ['seed']
+__all__ = ['seed', 'bernoulli']
 
 
 def seed(seed, ctx='all'):  # pylint: disable=redefined-outer-name
@@ -72,3 +73,10 @@ def seed(seed, ctx='all'):  # pylint: disable=redefined-outer-name
     array(0.9894903, ctx=gpu(0))
     """
     _mx_rand.seed(seed_state=seed, ctx=ctx)
+
+
+def bernoulli(probs=None, logits=None, size=None, dtype=None, ctx=None, out=None):
+    """
+    Sampling from bernoulli distribution.
+    """
+    return _mx_nd_npx.random.bernoulli(probs, logits, size, dtype, ctx, out)

python/mxnet/symbol/numpy_extension/__init__.py

@@ -19,6 +19,7 @@
 
 from . import _op
 from . import image
+from . import random
 from . import _register
 from ._op import *  # pylint: disable=wildcard-import
 

python/mxnet/symbol/numpy_extension/random.py

@@ -0,0 +1,57 @@
+# 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.
+
+"""Namespace for operators used in Gluon dispatched by F=symbol."""
+
+from __future__ import absolute_import
+from ...context import current_context
+from .. import _internal as _npi
+
+__all__ = ['bernoulli']
+
+
+def bernoulli(probs=None, logits=None, size=None, dtype=None, ctx=None, out=None):
+    """
+    Sampling from beroulli distributions.
+    """
+    from ..numpy import _Symbol as np_symbol
+    tensor_type_name = np_symbol
+    if (probs is None) == (logits is None):
+        raise ValueError(
+            "Either `probs` or `logits` must be specified, but not both.")
+    if dtype is None:
+        dtype = 'float32'
+    if ctx is None:
+        ctx = current_context()
+    if size == ():
+        size = None
+    if probs is not None:
+        is_tensor = isinstance(probs, tensor_type_name)
+        if is_tensor:
+            return _npi.bernoulli(probs, probs=None, logits=None, is_logit=False,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)
+        else:
+            return _npi.bernoulli(probs=probs, logits=None, is_logit=False,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)
+    else:
+        is_tensor = isinstance(logits, tensor_type_name)
+        if is_tensor:
+            return _npi.bernoulli(logits, probs=None, logits=None, is_logit=True,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)
+        else:
+            return _npi.bernoulli(probs=None, logits=logits, is_logit=True,
+                                  size=size, ctx=ctx, dtype=dtype, out=out)

src/operator/numpy/random/dist_common.h

@@ -27,7 +27,6 @@
 #ifndef MXNET_OPERATOR_NUMPY_RANDOM_DIST_COMMON_H_
 #define MXNET_OPERATOR_NUMPY_RANDOM_DIST_COMMON_H_
 
-#include <mshadow/base.h>
 #include <mxnet/operator_util.h>
 #include <algorithm>
 #include <string>
@@ -172,10 +171,36 @@ inline bool TwoparamsDistOpShape(const nnvm::NodeAttrs &attrs,
     } else if (in_attrs->size() == 0) {
       // Two scalar case.
       SHAPE_ASSIGN_CHECK(*out_attrs, 0, TShape(0, -1))
-      return true;
     }
   }
-  return out_attrs->at(0).ndim() != 0U;
+  return shape_is_known(out_attrs->at(0));
+}
+
+template <typename DistParam>
+inline bool UnaryDistOpShape(const nnvm::NodeAttrs &attrs,
+                                 std::vector<TShape> *in_attrs,
+                                 std::vector<TShape> *out_attrs) {
+  const DistParam &param = nnvm::get<DistParam>(attrs.parsed);
+  if (param.size.has_value()) {
+    // Size declared.
+    std::vector<dim_t> oshape_vec;
+    const mxnet::Tuple<int> &size = param.size.value();
+    for (int i = 0; i < size.ndim(); ++i) {
+      oshape_vec.emplace_back(size[i]);
+    }
+    SHAPE_ASSIGN_CHECK(*out_attrs, 0, TShape(oshape_vec));
+    for (size_t input_idx = 0; input_idx < in_attrs->size(); input_idx++) {
+      CheckBroadcastable((*in_attrs)[input_idx], (*out_attrs)[0]);
+    }
+  } else {
+    if (in_attrs->size() == 1U) {
+      // One param from ndarray.
+      SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0))
+    } else {
+      SHAPE_ASSIGN_CHECK(*out_attrs, 0, TShape(0, -1))
+    }
+  }
+  return shape_is_known(out_attrs->at(0));
 }
 
 }  // namespace op

src/operator/numpy/random/np_bernoulli_op.cc

@@ -0,0 +1,71 @@
+/*
+ * 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_bernoulli_op.cc
+ * \brief Operator for numpy sampling from bernoulli distributions
+ */
+
+#include "./np_bernoulli_op.h"
+#include "./dist_common.h"
+
+namespace mxnet {
+namespace op {
+
+DMLC_REGISTER_PARAMETER(NumpyBernoulliParam);
+
+NNVM_REGISTER_OP(_npi_bernoulli)
+.describe("Sample frmo bernoulli distribution")
+.set_num_inputs(
+  [](const nnvm::NodeAttrs& attrs) {
+    const NumpyBernoulliParam& param = nnvm::get<NumpyBernoulliParam>(attrs.parsed);
+    int num_inputs = 1;
+    if (param.logit.has_value() || param.prob.has_value()) {
+      num_inputs -= 1;
+    }
+    return num_inputs;
+  }
+)
+.set_num_outputs(1)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+  [](const NodeAttrs& attrs) {
+    const NumpyBernoulliParam& param = nnvm::get<NumpyBernoulliParam>(attrs.parsed);
+    int num_inputs = 1;
+    if (param.logit.has_value() || param.prob.has_value()) {
+      num_inputs -= 1;
+    }
+    if (num_inputs == 0) return std::vector<std::string>();
+    return std::vector<std::string>{"input1"};
+  })
+.set_attr_parser(ParamParser<NumpyBernoulliParam>)
+.set_attr<mxnet::FInferShape>("FInferShape", UnaryDistOpShape<NumpyBernoulliParam>)
+.set_attr<nnvm::FInferType>("FInferType", NumpyBernoulliOpType)
+.set_attr<FResourceRequest>("FResourceRequest",
+  [](const nnvm::NodeAttrs& attrs) {
+      return std::vector<ResourceRequest>{
+        ResourceRequest::kRandom, ResourceRequest::kTempSpace};
+  })
+.set_attr<FCompute>("FCompute<cpu>", NumpyBernoulliForward<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
+.add_argument("input1", "NDArray-or-Symbol", "Source input")
+.add_arguments(NumpyBernoulliParam::__FIELDS__());
+
+}  // namespace op
+}  // namespace mxnet

src/operator/numpy/random/np_bernoulli_op.cu

@@ -0,0 +1,35 @@
+/*
+ * 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_bernoulli_op.cu
+ * \brief Operator for numpy sampling from bernoulli distributions
+ */
+
+#include "./np_bernoulli_op.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_npi_bernoulli)
+.set_attr<FCompute>("FCompute<gpu>", NumpyBernoulliForward<gpu>);
+
+}  // namespace op
+}  // namespace mxnet