From 1dc93053290e0c42f2d5f5515ba6a3577b74dd3e Mon Sep 17 00:00:00 2001 From: Brian Patton Date: Wed, 5 Feb 2020 14:04:49 -0800 Subject: [PATCH] Adds a `tf.random.stateless_binomial` (stateless analogue to tf.random.Generator.binomial). PiperOrigin-RevId: 293446442 Change-Id: I1fc995304adbf82a5f43fddd27d015b6133b5a31 --- .../api_def_StatelessRandomBinomial.pbtxt | 48 +++++++ tensorflow/core/kernels/BUILD | 1 + tensorflow/core/kernels/random_binomial_op.cc | 118 ++++++++++++++++-- tensorflow/core/ops/stateless_random_ops.cc | 12 ++ .../eager/pywrap_gradient_exclusions.cc | 2 + .../random/random_binomial_test.py | 13 ++ tensorflow/python/ops/stateless_random_ops.py | 69 ++++++++++ .../api/golden/v1/tensorflow.random.pbtxt | 4 + .../api/golden/v1/tensorflow.raw_ops.pbtxt | 4 + .../api/golden/v2/tensorflow.random.pbtxt | 4 + .../api/golden/v2/tensorflow.raw_ops.pbtxt | 4 + 11 files changed, 268 insertions(+), 11 deletions(-) create mode 100644 tensorflow/core/api_def/base_api/api_def_StatelessRandomBinomial.pbtxt diff --git a/tensorflow/core/api_def/base_api/api_def_StatelessRandomBinomial.pbtxt b/tensorflow/core/api_def/base_api/api_def_StatelessRandomBinomial.pbtxt new file mode 100644 index 00000000000..3a7fbc9910b --- /dev/null +++ b/tensorflow/core/api_def/base_api/api_def_StatelessRandomBinomial.pbtxt @@ -0,0 +1,48 @@ +op { + graph_op_name: "StatelessRandomBinomial" + visibility: HIDDEN + in_arg { + name: "shape" + description: < +class StatelessRandomBinomialOp : public OpKernel { + // Reshape batches so each batch is this size if possible. + static const int32 kDesiredBatchSize = 100; + + public: + explicit StatelessRandomBinomialOp(OpKernelConstruction* context) + : OpKernel(context) {} + + void Compute(OpKernelContext* ctx) override { + const Tensor& shape_tensor = ctx->input(0); + const Tensor& seed_tensor = ctx->input(1); + const Tensor& counts_tensor = ctx->input(2); + const Tensor& probs_tensor = ctx->input(3); + + OP_REQUIRES(ctx, seed_tensor.dims() == 1 && seed_tensor.dim_size(0) == 2, + errors::InvalidArgument("seed must have shape [2], not ", + seed_tensor.shape().DebugString())); + + tensorflow::BCast bcast(counts_tensor.shape().dim_sizes(), + probs_tensor.shape().dim_sizes(), + /*fewer_dims_optimization=*/false, + /*return_flattened_batch_indices=*/true); + OP_REQUIRES(ctx, bcast.IsValid(), + errors::InvalidArgument( + "counts and probs must have compatible batch dimensions: ", + counts_tensor.shape().DebugString(), " vs. ", + probs_tensor.shape().DebugString())); + OP_REQUIRES( + ctx, TensorShapeUtils::IsVector(shape_tensor.shape()), + errors::InvalidArgument("Input shape should be a vector, got shape: ", + shape_tensor.shape().DebugString())); + OP_REQUIRES(ctx, + (shape_tensor.dtype() == DataType::DT_INT32 || + shape_tensor.dtype() == DataType::DT_INT64), + errors::InvalidArgument( + "Input shape should have dtype {int32, int64}.")); + + // Let's check that the shape tensor dominates the broadcasted tensor. + TensorShape bcast_shape = BCast::ToShape(bcast.output_shape()); + TensorShape output_shape; + if (shape_tensor.dtype() == DataType::DT_INT32) { + OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(shape_tensor.vec(), + &output_shape)); + } else { + OP_REQUIRES_OK(ctx, TensorShapeUtils::MakeShape(shape_tensor.vec(), + &output_shape)); + } + OP_REQUIRES(ctx, TensorShapeUtils::EndsWith(output_shape, bcast_shape), + errors::InvalidArgument( + "Shape passed in must end with broadcasted shape.")); + // Now that we have a guarantee, we can get the additional dimensions added + // by sampling. + int64 samples_per_batch = 1; + const int64 num_sample_dims = + (shape_tensor.dim_size(0) - bcast.output_shape().size()); + for (int64 i = 0; i < num_sample_dims; ++i) { + samples_per_batch *= shape_tensor.flat()(i); + } + int64 num_batches = 1; + for (int64 i = num_sample_dims; i < shape_tensor.dim_size(0); ++i) { + num_batches *= shape_tensor.flat()(i); + } + const int64 num_elements = num_batches * samples_per_batch; + + Tensor* samples_tensor; + OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &samples_tensor)); + if (output_shape.num_elements() == 0) return; + + random::PhiloxRandom::Key key; + random::PhiloxRandom::ResultType counter; + OP_REQUIRES_OK(ctx, GenerateKey(seed_tensor, &key, &counter)); + + auto philox = random::PhiloxRandom(counter, key); + auto binomial_functor = functor::RandomBinomialFunctor(); + binomial_functor(ctx, ctx->eigen_device(), num_batches, + samples_per_batch, num_elements, bcast, + counts_tensor.flat(), probs_tensor.flat(), philox, + samples_tensor->flat()); + } + + private: + TF_DISALLOW_COPY_AND_ASSIGN(StatelessRandomBinomialOp); +}; + } // namespace -#define REGISTER(RTYPE, TYPE) \ - REGISTER_KERNEL_BUILDER(Name("StatefulRandomBinomial") \ - .Device(DEVICE_CPU) \ - .HostMemory("resource") \ - .HostMemory("algorithm") \ - .HostMemory("shape") \ - .HostMemory("counts") \ - .HostMemory("probs") \ - .TypeConstraint("dtype") \ - .TypeConstraint("T"), \ - RandomBinomialOp) +#define REGISTER(RTYPE, TYPE) \ + REGISTER_KERNEL_BUILDER(Name("StatefulRandomBinomial") \ + .Device(DEVICE_CPU) \ + .HostMemory("resource") \ + .HostMemory("algorithm") \ + .HostMemory("shape") \ + .HostMemory("counts") \ + .HostMemory("probs") \ + .TypeConstraint("dtype") \ + .TypeConstraint("T"), \ + RandomBinomialOp); \ + REGISTER_KERNEL_BUILDER(Name("StatelessRandomBinomial") \ + .Device(DEVICE_CPU) \ + .HostMemory("shape") \ + .HostMemory("seed") \ + .HostMemory("counts") \ + .HostMemory("probs") \ + .TypeConstraint("dtype") \ + .TypeConstraint("T"), \ + StatelessRandomBinomialOp) #define REGISTER_ALL(RTYPE) \ REGISTER(RTYPE, Eigen::half); \ diff --git a/tensorflow/core/ops/stateless_random_ops.cc b/tensorflow/core/ops/stateless_random_ops.cc index f919a21d607..28b8271c5a0 100644 --- a/tensorflow/core/ops/stateless_random_ops.cc +++ b/tensorflow/core/ops/stateless_random_ops.cc @@ -93,4 +93,16 @@ REGISTER_OP("StatelessMultinomial") return Status::OK(); }); +REGISTER_OP("StatelessRandomBinomial") + .Input("shape: S") + .Input("seed: Tseed") + .Input("counts: T") + .Input("probs: T") + .Output("output: dtype") + .Attr("S: {int32, int64}") + .Attr("Tseed: {int32, int64} = DT_INT64") + .Attr("T: {half, float, double, int32, int64} = DT_DOUBLE") + .Attr("dtype: {half, float, double, int32, int64} = DT_INT64") + .SetShapeFn(StatelessShape); + } // namespace tensorflow diff --git a/tensorflow/python/eager/pywrap_gradient_exclusions.cc b/tensorflow/python/eager/pywrap_gradient_exclusions.cc index 80c066da88b..fb84ca207fc 100644 --- a/tensorflow/python/eager/pywrap_gradient_exclusions.cc +++ b/tensorflow/python/eager/pywrap_gradient_exclusions.cc @@ -307,6 +307,7 @@ bool OpGradientDoesntRequireInputIndices( {"StackPop", {true, {}}}, {"StackPush", {true, {}}}, {"StatelessMultinomial", {true, {}}}, + {"StatelessRandomBinomial", {true, {}}}, {"StatelessRandomNormal", {true, {}}}, {"StatelessRandomUniform", {true, {}}}, {"StatelessRandomUniformInt", {true, {}}}, @@ -761,6 +762,7 @@ bool OpGradientDoesntRequireOutputIndices( {"StackPop", {true, {}}}, {"StackPush", {true, {}}}, {"StatelessMultinomial", {true, {}}}, + {"StatelessRandomBinomial", {true, {}}}, {"StatelessRandomNormal", {true, {}}}, {"StatelessRandomUniform", {true, {}}}, {"StatelessRandomUniformInt", {true, {}}}, diff --git a/tensorflow/python/kernel_tests/random/random_binomial_test.py b/tensorflow/python/kernel_tests/random/random_binomial_test.py index 11bfd149c3f..a3ead328442 100644 --- a/tensorflow/python/kernel_tests/random/random_binomial_test.py +++ b/tensorflow/python/kernel_tests/random/random_binomial_test.py @@ -24,6 +24,7 @@ from tensorflow.python.framework import test_util from tensorflow.python.kernel_tests.random import util from tensorflow.python.ops import array_ops from tensorflow.python.ops import stateful_random_ops +from tensorflow.python.ops import stateless_random_ops from tensorflow.python.platform import test from tensorflow.python.platform import tf_logging @@ -80,6 +81,18 @@ class RandomBinomialTest(test.TestCase): sy = self._Sampler(1000, counts=10., probs=0.4, dtype=dt, seed=345) self.assertAllEqual(self.evaluate(sx()), self.evaluate(sy())) + def testStateless(self): + for dt in dtypes.float16, dtypes.float32, dtypes.float64: + sx = stateless_random_ops.stateless_random_binomial( + shape=[1000], seed=[12, 34], counts=10., probs=0.4, output_dtype=dt) + sy = stateless_random_ops.stateless_random_binomial( + shape=[1000], seed=[12, 34], counts=10., probs=0.4, output_dtype=dt) + sx0, sx1 = self.evaluate(sx), self.evaluate(sx) + sy0, sy1 = self.evaluate(sy), self.evaluate(sy) + self.assertAllEqual(sx0, sx1) + self.assertAllEqual(sx0, sy0) + self.assertAllEqual(sy0, sy1) + def testZeroShape(self): rnd = stateful_random_ops.Generator.from_seed(12345).binomial([0], [], []) self.assertEqual([0], rnd.shape.as_list()) diff --git a/tensorflow/python/ops/stateless_random_ops.py b/tensorflow/python/ops/stateless_random_ops.py index 62f65ced8d6..5ac687f5847 100644 --- a/tensorflow/python/ops/stateless_random_ops.py +++ b/tensorflow/python/ops/stateless_random_ops.py @@ -27,6 +27,7 @@ from tensorflow.python.util import deprecation from tensorflow.python.util.tf_export import tf_export ops.NotDifferentiable("StatelessMultinomial") +ops.NotDifferentiable("StatelessRandomBinomial") ops.NotDifferentiable("StatelessRandomNormal") ops.NotDifferentiable("StatelessRandomUniform") ops.NotDifferentiable("StatelessRandomUniformInt") @@ -102,6 +103,74 @@ def stateless_random_uniform(shape, return result +@tf_export("random.stateless_binomial") +def stateless_random_binomial(shape, + seed, + counts, + probs, + output_dtype=dtypes.int32, + name=None): + """Outputs deterministic pseudorandom values from a binomial distribution. + + The generated values follow a binomial distribution with specified count and + probability of success parameters. + + This is a stateless version of `tf.random.Generator.binomial`: if run twice + with the same seeds, it will produce the same pseudorandom numbers. The + output is consistent across multiple runs on the same hardware (and between + CPU and GPU), but may change between versions of TensorFlow or on non-CPU/GPU + hardware. + + Example: + + ```python + counts = [10., 20.] + # Probability of success. + probs = [0.8] + + binomial_samples = tf.random.stateless_binomial( + shape=[2], seed=[123, 456], counts=counts, probs=probs) + + counts = ... # Shape [3, 1, 2] + probs = ... # Shape [1, 4, 2] + shape = [3, 4, 3, 4, 2] + # Sample shape will be [3, 4, 3, 4, 2] + binomial_samples = tf.random.stateless_binomial( + shape=shape, seed=[123, 456], counts=counts, probs=probs) + ``` + + Args: + shape: A 1-D integer Tensor or Python array. The shape of the output tensor. + seed: A shape [2] integer Tensor of seeds to the random number generator. + counts: Tensor. The counts of the binomial distribution. Must be + broadcastable with `probs`, and broadcastable with the rightmost + dimensions of `shape`. + probs: Tensor. The probability of success for the binomial distribution. + Must be broadcastable with `counts` and broadcastable with the rightmost + dimensions of `shape`. + output_dtype: The type of the output. Default: tf.int32 + name: A name for the operation (optional). + + Returns: + samples: A Tensor of the specified shape filled with random binomial + values. For each i, each samples[..., i] is an independent draw from + the binomial distribution on counts[i] trials with probability of + success probs[i]. + + """ + with ops.name_scope(name, "stateless_random_binomial", + [shape, seed, counts, probs]) as name: + shape = tensor_util.shape_tensor(shape) + probs = ops.convert_to_tensor( + probs, dtype_hint=dtypes.float32, name="probs") + counts = ops.convert_to_tensor( + counts, dtype_hint=probs.dtype, name="counts") + result = gen_stateless_random_ops.stateless_random_binomial( + shape=shape, seed=seed, counts=counts, probs=probs, dtype=output_dtype) + tensor_util.maybe_set_static_shape(result, shape) + return result + + @tf_export("random.stateless_normal") def stateless_random_normal(shape, seed, diff --git a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt index 6d585aebd9a..6d01a3468ee 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.random.pbtxt @@ -72,6 +72,10 @@ tf_module { name: "shuffle" argspec: "args=[\'value\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } + member_method { + name: "stateless_binomial" + argspec: "args=[\'shape\', \'seed\', \'counts\', \'probs\', \'output_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " + } member_method { name: "stateless_categorical" argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " diff --git a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt index d92173083dc..630b78f15ef 100644 --- a/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt +++ b/tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt @@ -4280,6 +4280,10 @@ tf_module { name: "StatelessMultinomial" argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'output_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " } + member_method { + name: "StatelessRandomBinomial" + argspec: "args=[\'shape\', \'seed\', \'counts\', \'probs\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " + } member_method { name: "StatelessRandomNormal" argspec: "args=[\'shape\', \'seed\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " diff --git a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt index ae0da656caf..ff2b8ca18ce 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.random.pbtxt @@ -64,6 +64,10 @@ tf_module { name: "shuffle" argspec: "args=[\'value\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], " } + member_method { + name: "stateless_binomial" + argspec: "args=[\'shape\', \'seed\', \'counts\', \'probs\', \'output_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " + } member_method { name: "stateless_categorical" argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " diff --git a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt index d92173083dc..630b78f15ef 100644 --- a/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt +++ b/tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt @@ -4280,6 +4280,10 @@ tf_module { name: "StatelessMultinomial" argspec: "args=[\'logits\', \'num_samples\', \'seed\', \'output_dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " } + member_method { + name: "StatelessRandomBinomial" + argspec: "args=[\'shape\', \'seed\', \'counts\', \'probs\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], " + } member_method { name: "StatelessRandomNormal" argspec: "args=[\'shape\', \'seed\', \'dtype\', \'name\'], varargs=None, keywords=None, defaults=[\"\", \'None\'], "