Vulnerability: Heap out of bounds read in `MaxPoolGradWithArgmax` ### Impact
The implementation of `tf.raw_ops.MaxPoolGradWithArgmax` can cause reads outside of bounds of heap allocated data if attacker supplies specially crafted inputs:

```python
import tensorflow as tf

input = tf.constant([10.0, 10.0, 10.0], shape=[1, 1, 3, 1], dtype=tf.float32)
grad = tf.constant([10.0, 10.0, 10.0, 10.0], shape=[1, 1, 1, 4], dtype=tf.float32)
argmax = tf.constant([1], shape=[1], dtype=tf.int64)
ksize = [1, 1, 1, 1]
strides = [1, 1, 1, 1]
  
tf.raw_ops.MaxPoolGradWithArgmax(
  input=input, grad=grad, argmax=argmax, ksize=ksize, strides=strides,
  padding='SAME', include_batch_in_index=False)
```

The [implementation](https://github.com/tensorflow/tensorflow/blob/ef0c008ee84bad91ec6725ddc42091e19a30cf0e/tensorflow/core/kernels/maxpooling_op.cc#L1016-L1017) uses the same value to index in two different arrays but there is no guarantee that the sizes are identical. 

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

The fix will be included in TensorFlow 2.5.0. We will also cherrypick this commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow 2.1.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 Ying Wang and Yakun Zhang of Baidu X-Team.
Vulnerability ID: GHSA-545v-42p7-98fq
CWE: CWE-125
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/maxpooling_op.cc
Commit message: Fix heap buffer overflow

PiperOrigin-RevId: 372132844
Change-Id: Idef9895efaf145f2b1c23d31983601ec980cd5e4
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 struct LaunchMaxPoolingGradWithArgmax<CPUDevice, T> {         const int input_start = start * input_size_per_batch;         const int input_end = limit * input_size_per_batch;         for (int64 index = input_start; index < input_end; index++) {           int64 grad_out_index = argmax_flat(index);           if (!include_batch_in_index) {             const int64 cur_batch = index / input_size_per_batch;

Fixed version (after fix):
 struct LaunchMaxPoolingGradWithArgmax<CPUDevice, T> {         const int input_start = start * input_size_per_batch;         const int input_end = limit * input_size_per_batch;         for (int64 index = input_start; index < input_end; index++) {           if (index >= argmax.NumElements()) {             break;           }           int64 grad_out_index = argmax_flat(index);           if (!include_batch_in_index) {             const int64 cur_batch = index / input_size_per_batch;

Patch:
patch: @@ -1014,6 +1014,9 @@ struct LaunchMaxPoolingGradWithArgmax<CPUDevice, T> {
         const int input_start = start * input_size_per_batch;
         const int input_end = limit * input_size_per_batch;
         for (int64 index = input_start; index < input_end; index++) {
+          if (index >= argmax.NumElements()) {
+            break;
+          }
           int64 grad_out_index = argmax_flat(index);
           if (!include_batch_in_index) {
             const int64 cur_batch = index / input_size_per_batch;
path: tensorflow/core/kernels/maxpooling_op.cc