Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent out-of-bound accesses in SparseBincount.
PiperOrigin-RevId: 399918616
Change-Id: I11d154f4444d3fde1f09c5c40628b8671791a30d
  • Loading branch information
penpornk authored and tensorflower-gardener committed Sep 30, 2021
1 parent 4656caa commit f410212
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions tensorflow/core/kernels/bincount_op.cc
Expand Up @@ -405,6 +405,16 @@ class SparseBincountOp : public OpKernel {
for (int64_t i = 0; i < indices_mat.dimension(0); ++i) {
const int64_t batch = indices_mat(i, 0);
const Tidx bin = values(i);
OP_REQUIRES(
ctx, batch < out.dimension(0),
errors::InvalidArgument("Index out of bound. `batch` (", batch,
") must be less than the dimension size (",
out.dimension(0), ")."));
OP_REQUIRES(
ctx, bin < out.dimension(1),
errors::InvalidArgument("Index out ouf bound. `bin` (", bin,
") must be less then the dimension size (",
out.dimension(1), ")."));
if (bin < size) {
if (binary_output_) {
out(batch, bin) = T(1);
Expand Down

0 comments on commit f410212

Please sign in to comment.