Implements the sampling of truncated normal distribution using the inversed cumulative distribution function (CDF) method as described in https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf.
| 40 | // inversed cumulative distribution function (CDF) method as described in |
| 41 | // https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf. |
| 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).ConsumeValueOrDie().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 |
no test coverage detected