1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/tf2xla/lib/random.h"
17
18 #include <cmath>
19 #include <limits>
20
21 #include "tensorflow/compiler/tf2xla/xla_helpers.h"
22 #include "tensorflow/compiler/xla/client/lib/constants.h"
23 #include "tensorflow/compiler/xla/client/lib/math.h"
24 #include "tensorflow/compiler/xla/client/xla_builder.h"
25 #include "tensorflow/compiler/xla/status_macros.h"
26
27 namespace tensorflow {
28
TruncatedNormal(xla::XlaOp uniform)29 xla::XlaOp TruncatedNormal(xla::XlaOp uniform) {
30 const double kA = -2.0;
31 const double kB = 2.0;
32 const double kMu = 0.0;
33 const double kSigma = 1.0;
34 return ParameterizedTruncatedNormal(
35 uniform, xla::ScalarLike(uniform, kMu), xla::ScalarLike(uniform, kSigma),
36 xla::ScalarLike(uniform, kA), xla::ScalarLike(uniform, kB));
37 }
38
39 // Implements the sampling of truncated normal distribution using the
40 // inversed cumulative distribution function (CDF) method as described in
41 // https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf.
ParameterizedTruncatedNormal(xla::XlaOp uniform,xla::XlaOp mu,xla::XlaOp sigma,xla::XlaOp a,xla::XlaOp b)42 xla::XlaOp ParameterizedTruncatedNormal(xla::XlaOp uniform, xla::XlaOp mu,
43 xla::XlaOp sigma, xla::XlaOp a,
44 xla::XlaOp b) {
45 xla::XlaOp one = xla::ScalarLike(uniform, 1.0);
46 xla::XlaOp two = xla::ScalarLike(uniform, 2.0);
47 xla::XlaOp sqrt_2 = xla::ScalarLike(uniform, std::sqrt(2.0));
48
49 auto normal_cdf = [&](xla::XlaOp x) {
50 return (one + xla::Erf(x / sqrt_2)) / two;
51 };
52
53 // Calculate the cumulative probabilities for the lower and upper bound, a and
54 // b.
55 xla::XlaOp alpha = (a - mu) / sigma;
56 xla::XlaOp beta = (b - mu) / sigma;
57 xla::XlaOp alpha_normal_cdf = normal_cdf(alpha);
58 xla::XlaOp beta_normal_cdf = normal_cdf(beta);
59
60 // Convert the random uniform value in range (0, 1) (uniform) to a value in
61 // range (alpha_normal_cdf, beta_normal_cdf) that represents the cumulative
62 // probability (p) of a value (x) in the truncated normal distribution.
63 xla::XlaOp p =
64 alpha_normal_cdf + (beta_normal_cdf - alpha_normal_cdf) * uniform;
65
66 // Calculate x using the inversed cumulative distribution function:
67 // x = inversed_cdf(mu, sigma; p) = mu + sigma * sqrt(2) * erfinv(2*p-1)
68 // Clamp the input of erfinv to (-1, 1) because 2*p-1 may produce +/-1 due to
69 // computation precision.
70 xla::XlaOp v = two * p - one;
71 xla::PrimitiveType primitive_type =
72 uniform.builder()->GetShape(uniform).value().element_type();
73 xla::XlaOp epsilon = xla::Epsilon(uniform.builder(), primitive_type);
74 v = xla::Clamp(-one + epsilon, v, one - epsilon);
75 xla::XlaOp x = mu + sigma * sqrt_2 * xla::ErfInv(v);
76
77 // The value of x may be out of the range of (a, b), this typically happens
78 // when the region of the truncated normal has a very small probability.
79 x = xla::Clamp(a, x, b);
80
81 return x;
82 }
83
84 } // namespace tensorflow
85