Skip to content
Permalink
Browse files Browse the repository at this point in the history
Prevent heap OOB in sparse reduction ops.
PiperOrigin-RevId: 387934524
Change-Id: I894aa30f1e454f09b471d565b4a325da49322c1a
  • Loading branch information
mihaimaruseac authored and tensorflower-gardener committed Jul 31, 2021
1 parent 9c7f40e commit 87158f4
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions tensorflow/core/kernels/sparse_reduce_op.cc
Expand Up @@ -219,7 +219,20 @@ class SparseReduceOp : public OpKernel {
sp.Reorder<T>(reduction.reorder_dims);
for (const auto &g : sp.group(reduction.group_by_dims)) {
Op::template Run<T>(ctx, reduced_val, g.template values<T>());
OP_REQUIRES(ctx,
output_strides.empty() ||
(g.group().size() == output_strides.size()),
errors::Internal(
"Expected group size and output_strides size to match",
", but got ", g.group().size(), " and ",
output_strides.size()));
const int64_t idx = CoordinatesToFlatIndex(g.group(), output_strides);
OP_REQUIRES(ctx,
idx >= 0 && idx < out_flat.size(),
errors::Internal(
"Obtained a write index of ", idx,
" which is outside of bounds of [0, ",
out_flat.size(), ")"));
out_flat(idx) = reduced_val();
VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")
<< "; idx: " << idx << "; group " << Op::Name() << ": "
Expand Down

0 comments on commit 87158f4

Please sign in to comment.