Vulnerability: Heap out of bounds access in sparse reduction operations ### Impact
The implementation of sparse reduction operations in TensorFlow can trigger accesses outside of bounds of heap allocated data:

```python
import tensorflow as tf

x = tf.SparseTensor(
      indices=[[773, 773, 773], [773, 773, 773]],
      values=[1, 1],
      dense_shape=[337, 337, 337])
tf.sparse.reduce_sum(x, 1)
```

The [implementation](https://github.com/tensorflow/tensorflow/blob/a1bc56203f21a5a4995311825ffaba7a670d7747/tensorflow/core/kernels/sparse_reduce_op.cc#L217-L228) fails to validate that each reduction group does not overflow and that each corresponding index does not point to outside the bounds of the input tensor.

### Patches
We have patched the issue in GitHub commit [87158f43f05f2720a374f3e6d22a7aaa3a33f750](https://github.com/tensorflow/tensorflow/commit/87158f43f05f2720a374f3e6d22a7aaa3a33f750). 

The fix will be included in TensorFlow 2.6.0. We will also cherrypick this commit on TensorFlow 2.5.1, TensorFlow 2.4.3, and TensorFlow 2.3.4, as these are also affected and still in supported range.

### For more information
Please consult [our security guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for more information regarding the security model and how to contact us with issues and questions.

### Attribution
This vulnerability has been reported by members of the Aivul Team from Qihoo 360.
Vulnerability ID: GHSA-cgfm-62j4-v4rf
CWE: CWE-125
Severity score: 7.3
Project: https://github.com/tensorflow/tensorflow
Published: 2021-08-25
File: tensorflow/core/kernels/sparse_reduce_op.cc
Commit message: Prevent heap OOB in sparse reduction ops.

PiperOrigin-RevId: 387934524
Change-Id: I894aa30f1e454f09b471d565b4a325da49322c1a
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 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>());       const int64_t idx = CoordinatesToFlatIndex(g.group(), output_strides);       out_flat(idx) = reduced_val();       VLOG(2) << "coords: " << absl::StrJoin(g.group(), ",")               << "; idx: " << idx << "; group " << Op::Name() << ": "

Fixed version (after fix):
 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() << ": "

Patch:
patch: @@ -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() << ": "
path: tensorflow/core/kernels/sparse_reduce_op.cc