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

Commit 4de1f4a

Browse files
author
David Seiler
committed
PDF operators for each distribution for which we have a random sampler (plus also the PDF of the Dirichlet). Supports probabilities and log-probabilities, as well as gradients.
1 parent 99e69e6 commit 4de1f4a

File tree

5 files changed

+1109
-7
lines changed

5 files changed

+1109
-7
lines changed

‎python/mxnet/contrib/amp/lists/symbol.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,23 @@
600600
'_sparse_elemwise_mul',
601601
'_sparse_elemwise_sub',
602602
'_sparse_sum',
603+
604+
'random_pdf_gamma',
605+
'random_pdf_exponential',
606+
'random_pdf_uniform',
607+
'random_pdf_negative_binomial',
608+
'random_pdf_generalized_negative_binomial',
609+
'random_pdf_dirichlet',
610+
'random_pdf_normal',
611+
'random_pdf_poisson',
612+
'_random_pdf_gamma',
613+
'_random_pdf_exponential',
614+
'_random_pdf_uniform',
615+
'_random_pdf_negative_binomial',
616+
'_random_pdf_generalized_negative_binomial',
617+
'_random_pdf_dirichlet',
618+
'_random_pdf_normal',
619+
'_random_pdf_poisson',
603620
]
604621

605622
LOSS_OUTPUT_FUNCTIONS = [

‎src/operator/random/pdf_op.cc

Lines changed: 319 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,319 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
/*!
21+
* Copyright (c) 2018 by Contributors
22+
* \file pdf_op.cc
23+
* \brief CPU-operators for computing the pdf of random distributions.
24+
*/
25+
26+
#include "./pdf_op.h"
27+
28+
namespace mxnet {
29+
namespace op {
30+
31+
DMLC_REGISTER_PARAMETER(PdfParam);
32+
33+
#define MXNET_OPERATOR_REGISTER_PDF(distr, pdffunc, num_parms, \
34+
parm_name_1, parm_name_2, \
35+
parm_desc_1, parm_desc_2, \
36+
description, vectorparms) \
37+
NNVM_REGISTER_OP(_random_pdf_##distr) \
38+
.add_alias("random_pdf_" #distr) \
39+
.describe(description()+std::string(ADD_FILELINE)) \
40+
.set_num_inputs(num_parms+1) \
41+
.set_num_outputs(1) \
42+
.set_attr_parser(ParamParser<PdfParam>) \
43+
.set_attr<nnvm::FListInputNames>("FListInputNames", \
44+
[](const NodeAttrs& attrs) { \
45+
std::vector<std::string> v = {"sample", parm_name_1, parm_name_2}; \
46+
v.resize(num_parms+1); \
47+
return v; \
48+
}) \
49+
.set_attr<mxnet::FInferShape>("FInferShape", PdfOpShape<vectorparms>) \
50+
.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<num_parms+1, 1>) \
51+
.set_attr<FCompute>("FCompute<cpu>", PdfOpForward<cpu, pdffunc, num_parms, vectorparms>) \
52+
.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseInOut{"_backward_pdf_" #distr}) \
53+
.add_argument("sample", "NDArray-or-Symbol", "Samples from the distributions.") \
54+
.add_argument(parm_name_1, "NDArray-or-Symbol", parm_desc_1) \
55+
.add_arguments(PdfParam::__FIELDS__())
56+
57+
#define MXNET_OPERATOR_REGISTER_PDF_GRAD(distr, pdffunc, num_parms, vectorparms) \
58+
NNVM_REGISTER_OP(_backward_pdf_##distr) \
59+
.set_num_inputs(num_parms+3) \
60+
.set_num_outputs(num_parms+1) \
61+
.set_attr_parser(ParamParser<PdfParam>) \
62+
.set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs) \
63+
{ std::vector<std::pair<int, int> > v = {{1, 0}, {2, 1}, {3, 2}}; \
64+
v.resize(num_parms+1); \
65+
return v; }) \
66+
.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& attrs) \
67+
{ return std::vector<ResourceRequest>{ResourceRequest::kTempSpace}; }) \
68+
.set_attr<nnvm::TIsBackward>("TIsBackward", true) \
69+
.set_attr<FCompute>("FCompute<cpu>", PdfOpBackward<cpu, pdffunc##_Grad, num_parms, vectorparms>);
70+
71+
72+
#define MXNET_OPERATOR_REGISTER_PDF1(distr, pdffunc, parm_name, parm_desc, \
73+
description, vectorparms) \
74+
MXNET_OPERATOR_REGISTER_PDF(distr, pdffunc, 1, parm_name, parm_name, \
75+
parm_desc, parm_desc, description, vectorparms); \
76+
MXNET_OPERATOR_REGISTER_PDF_GRAD(distr, pdffunc, 1, vectorparms)
77+
78+
#define MXNET_OPERATOR_REGISTER_PDF2(distr, pdffunc, parm_name_1, parm_name_2, \
79+
parm_desc_1, parm_desc_2, description) \
80+
MXNET_OPERATOR_REGISTER_PDF(distr, pdffunc, 2, parm_name_1, parm_name_2, \
81+
parm_desc_1, parm_desc_2, description, false) \
82+
.add_argument(parm_name_2, "NDArray-or-Symbol", parm_desc_2); \
83+
MXNET_OPERATOR_REGISTER_PDF_GRAD(distr, pdffunc, 2, false)
84+
85+
inline std::string uniform_desc() {
86+
return std::string(R"code(Computes the value of the PDF of *sample* of
87+
uniform distributions on the intervals given by *[low,high)*.
88+
89+
*low* and *high* must have the same shape, which must match the leftmost subshape
90+
of *sample*. That is, *sample* can have the same shape as *low* and *high*, in which
91+
case the output contains one density per distribution, or *sample* can be a tensor
92+
of tensors with that shape, in which case the output is a tensor of densities such that
93+
the densities at index *i* in the output are given by the samples at index *i* in *sample*
94+
parameterized by the values of *low* and *high* at index *i*.
95+
96+
Examples::
97+
98+
random_pdf_uniform(sample=[[1,2,3,4]], low=[0], high=[10]) = [0.1, 0.1, 0.1, 0.1]
99+
100+
sample = [[[1, 2, 3],
101+
[1, 2, 3]],
102+
[[1, 2, 3],
103+
[1, 2, 3]]]
104+
low = [[0, 0],
105+
[0, 0]]
106+
high = [[ 5, 10],
107+
[15, 20]]
108+
random_pdf_uniform(sample=sample, low=low, high=high) =
109+
[[[0.2, 0.2, 0.2 ],
110+
[0.1, 0.1, 0.1 ]],
111+
[[0.06667, 0.06667, 0.06667],
112+
[0.05, 0.05, 0.05 ]]]
113+
114+
)code");
115+
}
116+
117+
inline std::string normal_desc() {
118+
return std::string(R"code(Computes the value of the PDF of *sample* of
119+
normal distributions with parameters *mu* (mean) and *sigma* (standard deviation).
120+
121+
*mu* and *sigma* must have the same shape, which must match the leftmost subshape
122+
of *sample*. That is, *sample* can have the same shape as *mu* and *sigma*, in which
123+
case the output contains one density per distribution, or *sample* can be a tensor
124+
of tensors with that shape, in which case the output is a tensor of densities such that
125+
the densities at index *i* in the output are given by the samples at index *i* in *sample*
126+
parameterized by the values of *mu* and *sigma* at index *i*.
127+
128+
Examples::
129+
130+
sample = [[-2, -1, 0, 1, 2]]
131+
random_pdf_normal(sample=sample, mu=[0], sigma=[1]) =
132+
[[0.05399097, 0.24197073, 0.3989423, 0.24197073, 0.05399097]]
133+
134+
random_pdf_normal(sample=sample*2, mu=[0,0], sigma=[1,2]) =
135+
[[0.05399097, 0.24197073, 0.3989423, 0.24197073, 0.05399097],
136+
[0.12098537, 0.17603266, 0.19947115, 0.17603266, 0.12098537]]
137+
)code");
138+
}
139+
140+
inline std::string gamma_desc() {
141+
return std::string(R"code(Computes the value of the PDF of *sample* of
142+
gamma distributions with parameters *alpha* (shape) and *beta* (rate).
143+
144+
*alpha* and *beta* must have the same shape, which must match the leftmost subshape
145+
of *sample*. That is, *sample* can have the same shape as *alpha* and *beta*, in which
146+
case the output contains one density per distribution, or *sample* can be a tensor
147+
of tensors with that shape, in which case the output is a tensor of densities such that
148+
the densities at index *i* in the output are given by the samples at index *i* in *sample*
149+
parameterized by the values of *alpha* and *beta* at index *i*.
150+
151+
Examples::
152+
153+
random_pdf_gamma(sample=[[1,2,3,4,5]], alpha=[5], beta=[1]) =
154+
[[0.01532831, 0.09022352, 0.16803136, 0.19536681, 0.17546739]]
155+
156+
sample = [[1, 2, 3, 4, 5],
157+
[2, 3, 4, 5, 6],
158+
[3, 4, 5, 6, 7]]
159+
160+
random_pdf_gamma(sample=sample, alpha=[5,6,7], beta=[1,1,1]) =
161+
[[0.01532831, 0.09022352, 0.16803136, 0.19536681, 0.17546739],
162+
[0.03608941, 0.10081882, 0.15629345, 0.17546739, 0.16062315],
163+
[0.05040941, 0.10419563, 0.14622283, 0.16062315, 0.14900276]]
164+
)code");
165+
}
166+
167+
inline std::string exponential_desc() {
168+
return std::string(R"code(Computes the value of the PDF of *sample* of
169+
exponential distributions with parameters *lam* (rate).
170+
171+
The shape of *lam* must match the leftmost subshape of *sample*. That is, *sample*
172+
can have the same shape as *lam*, in which case the output contains one density per
173+
distribution, or *sample* can be a tensor of tensors with that shape, in which case
174+
the output is a tensor of densities such that the densities at index *i* in the output
175+
are given by the samples at index *i* in *sample* parameterized by the value of *lam*
176+
at index *i*.
177+
178+
Examples::
179+
180+
random_pdf_exponential(sample=[[1, 2, 3]], lam=[1]) =
181+
[[0.36787945, 0.13533528, 0.04978707]]
182+
183+
sample = [[1,2,3],
184+
[1,2,3],
185+
[1,2,3]]
186+
187+
random_pdf_exponential(sample=sample, lam=[1,0.5,0.25]) =
188+
[[0.36787945, 0.13533528, 0.04978707],
189+
[0.30326533, 0.18393973, 0.11156508],
190+
[0.1947002, 0.15163267, 0.11809164]]
191+
)code");
192+
}
193+
194+
inline std::string poisson_desc() {
195+
return std::string(R"code(Computes the value of the PDF of *sample* of
196+
Poisson distributions with parameters *lam* (rate).
197+
198+
The shape of *lam* must match the leftmost subshape of *sample*. That is, *sample*
199+
can have the same shape as *lam*, in which case the output contains one density per
200+
distribution, or *sample* can be a tensor of tensors with that shape, in which case
201+
the output is a tensor of densities such that the densities at index *i* in the output
202+
are given by the samples at index *i* in *sample* parameterized by the value of *lam*
203+
at index *i*.
204+
205+
Examples::
206+
207+
random_pdf_poisson(sample=[[0,1,2,3]], lam=[1]) =
208+
[[0.36787945, 0.36787945, 0.18393973, 0.06131324]]
209+
210+
sample = [[0,1,2,3],
211+
[0,1,2,3],
212+
[0,1,2,3]]
213+
214+
random_pdf_poisson(sample=sample, lam=[1,2,3]) =
215+
[[0.36787945, 0.36787945, 0.18393973, 0.06131324],
216+
[0.13533528, 0.27067056, 0.27067056, 0.18044704],
217+
[0.04978707, 0.14936121, 0.22404182, 0.22404182]]
218+
)code");
219+
}
220+
221+
inline std::string negative_binomial_desc() {
222+
return std::string(R"code(Computes the value of the PDF of samples of
223+
negative binomial distributions with parameters *k* (failure limit) and *p* (failure probability).
224+
225+
*k* and *p* must have the same shape, which must match the leftmost subshape
226+
of *sample*. That is, *sample* can have the same shape as *k* and *p*, in which
227+
case the output contains one density per distribution, or *sample* can be a tensor
228+
of tensors with that shape, in which case the output is a tensor of densities such that
229+
the densities at index *i* in the output are given by the samples at index *i* in *sample*
230+
parameterized by the values of *k* and *p* at index *i*.
231+
232+
Examples::
233+
234+
random_pdf_negative_binomial(sample=[[1,2,3,4]], k=[1], p=a[0.5]) =
235+
[[0.25, 0.125, 0.0625, 0.03125]]
236+
237+
# Note that k may be real-valued
238+
sample = [[1,2,3,4],
239+
[1,2,3,4]]
240+
random_pdf_negative_binomial(sample=sample, k=[1, 1.5], p=[0.5, 0.5]) =
241+
[[0.25, 0.125, 0.0625, 0.03125 ],
242+
[0.26516506, 0.16572815, 0.09667476, 0.05437956]]
243+
)code");
244+
}
245+
246+
inline std::string generalized_negative_binomial_desc() {
247+
return std::string(R"code(Computes the value of the PDF of *sample* of
248+
generalized negative binomial distributions with parameters *mu* (mean)
249+
and *alpha* (dispersion). This can be understood as a reparameterization of
250+
the negative binomial, where *k* = *1 / alpha* and *p* = *1 / (mu \* alpha + 1)*.
251+
252+
*mu* and *alpha* must have the same shape, which must match the leftmost subshape
253+
of *sample*. That is, *sample* can have the same shape as *mu* and *alpha*, in which
254+
case the output contains one density per distribution, or *sample* can be a tensor
255+
of tensors with that shape, in which case the output is a tensor of densities such that
256+
the densities at index *i* in the output are given by the samples at index *i* in *sample*
257+
parameterized by the values of *mu* and *alpha* at index *i*.
258+
259+
Examples::
260+
261+
random_pdf_generalized_negative_binomial(sample=[[1, 2, 3, 4]], alpha=[1], mu=[1]) =
262+
[[0.25, 0.125, 0.0625, 0.03125]]
263+
264+
sample = [[1,2,3,4],
265+
[1,2,3,4]]
266+
random_pdf_generalized_negative_binomial(sample=sample, alpha=[1, 0.6666], mu=[1, 1.5]) =
267+
[[0.25, 0.125, 0.0625, 0.03125 ],
268+
[0.26517063, 0.16573331, 0.09667706, 0.05437994]]
269+
)code");
270+
}
271+
272+
inline std::string dirichlet_desc() {
273+
return std::string(R"code(Computes the value of the PDF of *sample* of
274+
Dirichlet distributions with parameter *alpha*.
275+
276+
The shape of *alpha* must match the leftmost subshape of *sample*. That is, *sample*
277+
can have the same shape as *alpha*, in which case the output contains one density per
278+
distribution, or *sample* can be a tensor of tensors with that shape, in which case
279+
the output is a tensor of densities such that the densities at index *i* in the output
280+
are given by the samples at index *i* in *sample* parameterized by the value of *alpha*
281+
at index *i*.
282+
283+
Examples::
284+
285+
random_pdf_dirichlet(sample=[[1,2],[2,3],[3,4]], alpha=[2.5, 2.5]) =
286+
[38.413498, 199.60245, 564.56085]
287+
288+
sample = [[[1, 2, 3], [10, 20, 30], [100, 200, 300]],
289+
[[0.1, 0.2, 0.3], [0.01, 0.02, 0.03], [0.001, 0.002, 0.003]]]
290+
291+
random_pdf_dirichlet(sample=sample, alpha=[0.1, 0.4, 0.9]) =
292+
[[2.3257459e-02, 5.8420084e-04, 1.4674458e-05],
293+
[9.2589635e-01, 3.6860607e+01, 1.4674468e+03]]
294+
)code");
295+
}
296+
297+
MXNET_OPERATOR_REGISTER_PDF2(uniform, PDF_Uniform, "low", "high",
298+
"Lower bounds of the distributions.", "Upper bounds of the distributions.", uniform_desc)
299+
MXNET_OPERATOR_REGISTER_PDF2(normal, PDF_Normal, "mu", "sigma",
300+
"Means of the distributions.", "Standard deviations of the distributions.", normal_desc)
301+
MXNET_OPERATOR_REGISTER_PDF2(gamma, PDF_Gamma, "alpha", "beta",
302+
"Alpha (shape) parameters of the distributions.", "Beta (scale) parameters of the distributions.",
303+
gamma_desc)
304+
MXNET_OPERATOR_REGISTER_PDF1(exponential, PDF_Exponential, "lam",
305+
"Lambda (rate) parameters of the distributions.", exponential_desc, false)
306+
MXNET_OPERATOR_REGISTER_PDF1(poisson, PDF_Poisson, "lam",
307+
"Lambda (rate) parameters of the distributions.", poisson_desc, false)
308+
MXNET_OPERATOR_REGISTER_PDF2(negative_binomial, PDF_NegativeBinomial, "k", "p",
309+
"Limits of unsuccessful experiments.", "Failure probabilities in each experiment.",
310+
negative_binomial_desc)
311+
MXNET_OPERATOR_REGISTER_PDF2(generalized_negative_binomial,
312+
PDF_GeneralizedNegativeBinomial, "mu", "alpha",
313+
"Means of the distributions.", "Alpha (dispersion) parameters of the distributions.",
314+
generalized_negative_binomial_desc)
315+
MXNET_OPERATOR_REGISTER_PDF1(dirichlet, PDF_Dirichlet, "alpha",
316+
"Concentration parameters of the distributions.", dirichlet_desc, true)
317+
318+
} // namespace op
319+
} // namespace mxnet

‎src/operator/random/pdf_op.cu

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
/*!
21+
* Copyright (c) 2018 by Contributors
22+
* \file pdf_op.cu
23+
* \brief GPU-operators for computing the pdf of random distributions.
24+
*/
25+
26+
#include "./pdf_op.h"
27+
28+
namespace mxnet {
29+
namespace op {
30+
31+
#define MXNET_OPERATOR_REGISTER_PDF(distr, pdffunc, num_parms, vector_parms) \
32+
NNVM_REGISTER_OP(_random_pdf_##distr) \
33+
.set_attr<FCompute>("FCompute<gpu>", PdfOpForward<gpu, pdffunc, num_parms, vector_parms>); \
34+
NNVM_REGISTER_OP(_backward_pdf_##distr) \
35+
.set_attr<FCompute>("FCompute<gpu>", PdfOpBackward<gpu, pdffunc##_Grad, num_parms, vector_parms>);
36+
37+
MXNET_OPERATOR_REGISTER_PDF(uniform, PDF_Uniform, 2, false)
38+
MXNET_OPERATOR_REGISTER_PDF(normal, PDF_Normal, 2, false)
39+
MXNET_OPERATOR_REGISTER_PDF(gamma, PDF_Gamma, 2, false)
40+
MXNET_OPERATOR_REGISTER_PDF(exponential, PDF_Exponential, 1, false)
41+
MXNET_OPERATOR_REGISTER_PDF(poisson, PDF_Poisson, 1, false)
42+
MXNET_OPERATOR_REGISTER_PDF(negative_binomial, PDF_NegativeBinomial, 2, false)
43+
MXNET_OPERATOR_REGISTER_PDF(generalized_negative_binomial,
44+
PDF_GeneralizedNegativeBinomial, 2, false)
45+
MXNET_OPERATOR_REGISTER_PDF(dirichlet, PDF_Dirichlet, 1, true)
46+
47+
} // namespace op
48+
} // namespace mxnet

0 commit comments

Comments
 (0)