Fix crash in range sampler by adding a range check in the sampler op.

Change: 147562709

1 files changed, 6 insertions, 2 deletions

diff --git a/tensorflow/core/kernels/candidate_sampler_ops.cc b/tensorflow/core/kernels/candidate_sampler_ops.cc
index 6aa9059dc7..9e8b122801 100644
--- a/tensorflow/core/kernels/candidate_sampler_ops.cc
+++ b/tensorflow/core/kernels/candidate_sampler_ops.cc
@@ -47,6 +47,12 @@ class BaseCandidateSamplerOp : public OpKernel {
     OP_REQUIRES(context, true_classes.dim_size(1) == num_true_,
                 errors::InvalidArgument("true_classes must have "
                                         "num_true columns"));
+    CHECK(sampler_) << "CandidateSamplerOp did not set sampler_";
+
+    if (unique_) {
+      OP_REQUIRES(context, num_sampled_ <= sampler_->range(),
+                  errors::InvalidArgument("Sampler's range is too small."));
+    }
 
     // Output candidates and expected_count.
     Tensor* out_sampled_candidates = nullptr;
@@ -73,8 +79,6 @@ class BaseCandidateSamplerOp : public OpKernel {
     gtl::MutableArraySlice<float> sampled_expected_count(
         out_sampled_expected_count->vec<float>().data(), num_sampled_);
 
-    CHECK(sampler_) << "CandidateSamplerOp did not set sampler_";
-
     // Approximately conservatively estimate the number of samples required.
     // In cases where rejection sampling is used we may occasionally use more
     // samples than expected, which will result in reused random bits.