1/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#ifndef TENSORFLOW_CORE_KERNELS_RANDOM_BINOMIAL_OP_H_
17#define TENSORFLOW_CORE_KERNELS_RANDOM_BINOMIAL_OP_H_
18
19#include "tensorflow/core/framework/tensor_types.h"
20#include "tensorflow/core/lib/random/random_distributions.h"
21
22namespace tensorflow {
23
24class OpKernelContext;
25
26namespace functor {
27
28// Sample a binomial random variable, with probs and counts for each batch.
29// Uses binomial inversion and a transformed rejection sampling method as
30// described in
31// https://pdfs.semanticscholar.org/471b/c2726e25bbf8801ef781630a2c13f654268e.pdf.
32// Two different algorithms are employed, depending on the size of
33// counts * probs (or counts * (1 - probs) if probs > 0.5.
34// If counts * probs < 10, we simply sum up Geometric random variables until
35// they exceed count, and the number we used is binomially distributed.
36// In expectation, this will take O(counts * probs) time, and requiring in
37// expectation the same number of random variates.
38// This can be much cheaper than summing bernoulli random variates, as we
39// will always need O(counts) bernoulli random variates (so this requires fewer
40// uniform r.v.s as well as can be faster).
41//
42// If counts * probs > 10, we use a transformed-rejection algorithm based on
43// pairs of uniform random variates due to Hormann.
44// https://pdfs.semanticscholar.org/471b/c2726e25bbf8801ef781630a2c13f654268e.pdf
45// This algorithm has higher acceptance rates for counts * probs large, as the
46// proposal distribution becomes quite tight, requiring approximately two
47// uniform random variates as counts * probs becomes large.
48template <typename Device, typename T, typename U>
49struct RandomBinomialFunctor {
50 void operator()(OpKernelContext* ctx, const Device& d, int64_t num_batches,
51 int64_t samples_per_batch, int64_t num_elements,
52 typename TTypes<T>::ConstFlat counts,
53 typename TTypes<T>::ConstFlat probs,
54 const random::PhiloxRandom& gen,
55 typename TTypes<U>::Flat output);
56};
57
58} // namespace functor
59} // namespace tensorflow
60
61#endif // TENSORFLOW_CORE_KERNELS_RANDOM_BINOMIAL_OP_H_
62