Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
sample_like operators (#13034)
Browse files Browse the repository at this point in the history
  • Loading branch information
@szha @eric-haibin-lin
szha authored and eric-haibin-lin committed Oct 30, 2018
1 parent a9ac9d1 commit afbb72f
Show file tree
Hide file tree
Showing 5 changed files with 669 additions and 291 deletions.
12 changes: 8 additions & 4 deletions python/mxnet/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -560,7 +560,7 @@ def _as_list(obj):
return [obj]


_OP_NAME_PREFIX_LIST = ['_contrib_', '_linalg_', '_sparse_', '_image_']
_OP_NAME_PREFIX_LIST = ['_contrib_', '_linalg_', '_sparse_', '_image_', '_random_']


def _get_op_name_prefix(op_name):
Expand Down Expand Up @@ -616,9 +616,13 @@ def _init_op_module(root_namespace, module_name, make_op_func):
op_name_prefix = _get_op_name_prefix(name)
module_name_local = module_name
if len(op_name_prefix) > 0:
func_name = name[len(op_name_prefix):]
cur_module = submodule_dict[op_name_prefix]
module_name_local = "%s.%s.%s" % (root_namespace, module_name, op_name_prefix[1:-1])
if op_name_prefix != '_random_' or name.endswith('_like'):
func_name = name[len(op_name_prefix):]
cur_module = submodule_dict[op_name_prefix]
module_name_local = "%s.%s.%s" % (root_namespace, module_name, op_name_prefix[1:-1])
else:
func_name = name
cur_module = module_internal
elif name.startswith('_'):
func_name = name
cur_module = module_internal
Expand Down
179 changes: 138 additions & 41 deletions src/operator/random/sample_op.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,15 +37,44 @@ DMLC_REGISTER_PARAMETER(SamplePoissonParam);
DMLC_REGISTER_PARAMETER(SampleNegBinomialParam);
DMLC_REGISTER_PARAMETER(SampleGenNegBinomialParam);

#define MXNET_OPERATOR_REGISTER_SAMPLE(name, ParamType) \
NNVM_REGISTER_OP(name) \
.set_num_inputs(0) \
.set_num_outputs(1) \
.set_attr_parser(ParamParser<ParamType>) \
.set_attr<nnvm::FInferShape>("FInferShape", InitShape<ParamType>) \
.set_attr<nnvm::FInferType>("FInferType", SampleOpType<ParamType>) \
.set_attr<FResourceRequest>("FResourceRequest", SampleResource) \
.add_arguments(ParamType::__FIELDS__())
DMLC_REGISTER_PARAMETER(SampleUniformLikeParam);
DMLC_REGISTER_PARAMETER(SampleNormalLikeParam);
DMLC_REGISTER_PARAMETER(SampleGammaLikeParam);
DMLC_REGISTER_PARAMETER(SampleExponentialLikeParam);
DMLC_REGISTER_PARAMETER(SamplePoissonLikeParam);
DMLC_REGISTER_PARAMETER(SampleNegBinomialLikeParam);
DMLC_REGISTER_PARAMETER(SampleGenNegBinomialLikeParam);

#define MXNET_OPERATOR_REGISTER_SAMPLE(name, ParamType) \
NNVM_REGISTER_OP(name) \
.set_num_inputs(0) \
.set_num_outputs(1) \
.set_attr_parser(ParamParser<ParamType>) \
.set_attr<nnvm::FInferShape>("FInferShape", InitShape<ParamType>) \
.set_attr<nnvm::FInferType>("FInferType", SampleOpType<ParamType>) \
.set_attr<FResourceRequest>("FResourceRequest", SampleResource) \
.add_arguments(ParamType::__FIELDS__()) \
.set_attr<FInferStorageType>("FInferStorageType", InitStorageType<ParamType, true, false>) \
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, ParamType>) \
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, ParamType>)

#define MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(name, ParamType) \
NNVM_REGISTER_OP(name) \
.set_num_inputs(1) \
.set_num_outputs(1) \
.set_attr_parser(ParamParser<ParamType>) \
.set_attr<nnvm::FInferShape>("FInferShape", ElemwiseShape<1, 1>) \
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>) \
.set_attr<FResourceRequest>("FResourceRequest", SampleResource) \
.set_attr<nnvm::FIgnoreInputs>("FIgnoreInputs", \
[](const NodeAttrs& attrs) { return std::vector<uint32_t>(1, 0); }) \
.set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes) \
.add_arguments(ParamType::__FIELDS__()) \
.add_argument("data", "NDArray-or-Symbol", "The input") \
.set_attr<FInferStorageType>("FInferStorageType", \
ElemwiseStorageType<1, 1, false, true, false>) \
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, ParamType>) \
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, ParamType>)

// Add "uniform" alias for backward compatibility
MXNET_OPERATOR_REGISTER_SAMPLE(_random_uniform, SampleUniformParam)
Expand All @@ -63,10 +92,7 @@ Example::
uniform(low=0, high=1, shape=(2,2)) = [[ 0.60276335, 0.85794562],
[ 0.54488319, 0.84725171]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType", InitStorageType<SampleUniformParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, UniformSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, UniformSampler<cpu>>);
)code" ADD_FILELINE);

// Add "normal" alias for backward compatibility
MXNET_OPERATOR_REGISTER_SAMPLE(_random_normal, SampleNormalParam)
Expand All @@ -76,16 +102,14 @@ MXNET_OPERATOR_REGISTER_SAMPLE(_random_normal, SampleNormalParam)
.. note:: The existing alias ``normal`` is deprecated.
Samples are distributed according to a normal distribution parametrized by *loc* (mean) and *scale* (standard deviation).
Samples are distributed according to a normal distribution parametrized by *loc* (mean) and *scale*
(standard deviation).
Example::
normal(loc=0, scale=1, shape=(2,2)) = [[ 1.89171135, -1.16881478],
[-1.23474145, 1.55807114]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType", InitStorageType<SampleNormalParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, NormalSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, NormalSampler<cpu>>);
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE(_random_gamma, SampleGammaParam)
.add_alias("random_gamma")
Expand All @@ -97,10 +121,7 @@ Example::
gamma(alpha=9, beta=0.5, shape=(2,2)) = [[ 7.10486984, 3.37695289],
[ 3.91697288, 3.65933681]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType", InitStorageType<SampleGammaParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, GammaSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, GammaSampler<cpu>>);
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE(_random_exponential, SampleExponentialParam)
.add_alias("random_exponential")
Expand All @@ -112,11 +133,7 @@ Example::
exponential(lam=4, shape=(2,2)) = [[ 0.0097189 , 0.08999364],
[ 0.04146638, 0.31715935]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType",
InitStorageType<SampleExponentialParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, ExponentialSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, ExponentialSampler<cpu>>);
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE(_random_poisson, SamplePoissonParam)
.add_alias("random_poisson")
Expand All @@ -129,10 +146,7 @@ Example::
poisson(lam=4, shape=(2,2)) = [[ 5., 2.],
[ 4., 6.]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType", InitStorageType<SamplePoissonParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, PoissonSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, PoissonSampler<cpu>>);
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE(_random_negative_binomial, SampleNegBinomialParam)
.add_alias("random_negative_binomial")
Expand All @@ -146,11 +160,7 @@ Example::
negative_binomial(k=3, p=0.4, shape=(2,2)) = [[ 4., 7.],
[ 2., 5.]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType",
InitStorageType<SampleNegBinomialParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, NegativeBinomialSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, NegativeBinomialSampler<cpu>>);
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE(_random_generalized_negative_binomial, SampleGenNegBinomialParam)
.add_alias("random_generalized_negative_binomial")
Expand All @@ -165,11 +175,98 @@ Example::
generalized_negative_binomial(mu=2.0, alpha=0.3, shape=(2,2)) = [[ 2., 1.],
[ 6., 4.]]
)code" ADD_FILELINE)
.set_attr<FInferStorageType>("FInferStorageType",
InitStorageType<SampleGenNegBinomialParam, true, false>)
.set_attr<FCompute>("FCompute<cpu>", Sample_<cpu, GeneralizedNegativeBinomialSampler<cpu>>)
.set_attr<FComputeEx>("FComputeEx<cpu>", SampleEx_<cpu, GeneralizedNegativeBinomialSampler<cpu>>);
)code" ADD_FILELINE);


// *_like operators

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_uniform_like, SampleUniformLikeParam)
.describe(R"code(Draw random samples from a uniform distribution according to the input array shape.
Samples are uniformly distributed over the half-open interval *[low, high)*
(includes *low*, but excludes *high*).
Example::
uniform(low=0, high=1, data=ones(2,2)) = [[ 0.60276335, 0.85794562],
[ 0.54488319, 0.84725171]]
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_normal_like, SampleNormalLikeParam)
.describe(R"code(Draw random samples from a normal (Gaussian) distribution according to the input array shape.
Samples are distributed according to a normal distribution parametrized by *loc* (mean) and *scale*
(standard deviation).
Example::
normal(loc=0, scale=1, data=ones(2,2)) = [[ 1.89171135, -1.16881478],
[-1.23474145, 1.55807114]]
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_gamma_like, SampleGammaLikeParam)
.describe(R"code(Draw random samples from a gamma distribution according to the input array shape.
Samples are distributed according to a gamma distribution parametrized by *alpha* (shape) and *beta* (scale).
Example::
gamma(alpha=9, beta=0.5, data=ones(2,2)) = [[ 7.10486984, 3.37695289],
[ 3.91697288, 3.65933681]]
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_exponential_like, SampleExponentialLikeParam)
.describe(R"code(Draw random samples from an exponential distribution according to the input array shape.
Samples are distributed according to an exponential distribution parametrized by *lambda* (rate).
Example::
exponential(lam=4, data=ones(2,2)) = [[ 0.0097189 , 0.08999364],
[ 0.04146638, 0.31715935]]
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_poisson_like, SamplePoissonLikeParam)
.describe(R"code(Draw random samples from a Poisson distribution according to the input array shape.
Samples are distributed according to a Poisson distribution parametrized by *lambda* (rate).
Samples will always be returned as a floating point data type.
Example::
poisson(lam=4, data=ones(2,2)) = [[ 5., 2.],
[ 4., 6.]]
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_negative_binomial_like, SampleNegBinomialLikeParam)
.describe(R"code(Draw random samples from a negative binomial distribution according to the input array shape.
Samples are distributed according to a negative binomial distribution parametrized by
*k* (limit of unsuccessful experiments) and *p* (failure probability in each experiment).
Samples will always be returned as a floating point data type.
Example::
negative_binomial(k=3, p=0.4, data=ones(2,2)) = [[ 4., 7.],
[ 2., 5.]]
)code" ADD_FILELINE);

MXNET_OPERATOR_REGISTER_SAMPLE_LIKE(_random_generalized_negative_binomial_like,
SampleGenNegBinomialLikeParam)
.describe(R"code(Draw random samples from a generalized negative binomial distribution according to the
input array shape.
Samples are distributed according to a generalized negative binomial distribution parametrized by
*mu* (mean) and *alpha* (dispersion). *alpha* is defined as *1/k* where *k* is the failure limit of the
number of unsuccessful experiments (generalized to real numbers).
Samples will always be returned as a floating point data type.
Example::
generalized_negative_binomial(mu=2.0, alpha=0.3, data=ones(2,2)) = [[ 2., 1.],
[ 6., 4.]]
)code" ADD_FILELINE);

} // namespace op
} // namespace mxnet
47 changes: 20 additions & 27 deletions src/operator/random/sample_op.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,33 +27,26 @@
namespace mxnet {
namespace op {

NNVM_REGISTER_OP(_random_uniform)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, UniformSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, UniformSampler<gpu>>);

NNVM_REGISTER_OP(_random_normal)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, NormalSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, NormalSampler<gpu>>);

NNVM_REGISTER_OP(_random_gamma)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, GammaSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, GammaSampler<gpu>>);

NNVM_REGISTER_OP(_random_exponential)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, ExponentialSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, ExponentialSampler<gpu>>);

NNVM_REGISTER_OP(_random_poisson)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, PoissonSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, PoissonSampler<gpu>>);

NNVM_REGISTER_OP(_random_negative_binomial)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, NegativeBinomialSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, NegativeBinomialSampler<gpu>>);

NNVM_REGISTER_OP(_random_generalized_negative_binomial)
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, GeneralizedNegativeBinomialSampler<gpu>>)
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, GeneralizedNegativeBinomialSampler<gpu>>);
#define MXNET_OPERATOR_REGISTER_SAMPLE_GPU(name, ParamType) \
NNVM_REGISTER_OP(name) \
.set_attr<FCompute>("FCompute<gpu>", Sample_<gpu, ParamType>) \
.set_attr<FComputeEx>("FComputeEx<gpu>", SampleEx_<gpu, ParamType>); \

MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_uniform, SampleUniformParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_normal, SampleNormalParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_gamma, SampleGammaParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_exponential, SampleExponentialParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_poisson, SamplePoissonParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_negative_binomial, SampleNegBinomialParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_generalized_negative_binomial, SampleGenNegBinomialParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_uniform_like, SampleUniformLikeParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_normal_like, SampleNormalLikeParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_gamma_like, SampleGammaLikeParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_exponential_like, SampleExponentialLikeParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_poisson_like, SamplePoissonLikeParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_negative_binomial_like, SampleNegBinomialLikeParam)
MXNET_OPERATOR_REGISTER_SAMPLE_GPU(_random_generalized_negative_binomial_like,
SampleGenNegBinomialLikeParam)

} // namespace op
} // namespace mxnet
Loading

0 comments on commit afbb72f

Please sign in to comment.