Vulnerability: Heap buffer overflow in `AvgPool3DGrad` ### Impact
The implementation of `tf.raw_ops.AvgPool3DGrad` is vulnerable to a heap buffer overflow:

```python
import tensorflow as tf

orig_input_shape = tf.constant([10, 6, 3, 7, 7], shape=[5], dtype=tf.int32)
grad = tf.constant([0.01, 0, 0], shape=[3, 1, 1, 1, 1], dtype=tf.float32)
ksize = [1, 1, 1, 1, 1]
strides = [1, 1, 1, 1, 1]
padding = "SAME"

tf.raw_ops.AvgPool3DGrad(
  orig_input_shape=orig_input_shape, grad=grad, ksize=ksize, strides=strides,
  padding=padding)
```

The [implementation](https://github.com/tensorflow/tensorflow/blob/d80ffba9702dc19d1fac74fc4b766b3fa1ee976b/tensorflow/core/kernels/pooling_ops_3d.cc#L376-L450) assumes that the `orig_input_shape` and `grad` tensors have similar first and last dimensions but does not check that this assumption is validated.

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

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-v6r6-84gr-92rm
CWE: CWE-787
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/pooling_ops_3d.cc
Commit message: Fix assertion failure in pooling_ops_3d

PiperOrigin-RevId: 372364504
Change-Id: Iecde4fe26b47a8fa935d6e2611b5585ed5777781
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 struct LaunchAvgPooling3dGradOp<CPUDevice, T> {                      const std::array<int64, 3>& output_shape,                      const std::array<int64, 3>& padding,                      TensorFormat data_format, Tensor* output) {     output->flat<T>().setZero();     std::array<int64, 3> input_size = {{tensor_in_shape.dim_size(3),                                         tensor_in_shape.dim_size(2),

Fixed version (after fix):
 struct LaunchAvgPooling3dGradOp<CPUDevice, T> {                      const std::array<int64, 3>& output_shape,                      const std::array<int64, 3>& padding,                      TensorFormat data_format, Tensor* output) {     OP_REQUIRES(         context, tensor_in_shape.dim_size(0) == out_backprop.dim_size(0),         errors::InvalidArgument(             "Expected first dimension of tensor_in_shape and "             "out_backprop to match, got ",             tensor_in_shape.dim_size(0), " and ", out_backprop.dim_size(0)));     OP_REQUIRES(         context, tensor_in_shape.dim_size(4) == out_backprop.dim_size(4),         errors::InvalidArgument(             "Expected last dimension of tensor_in_shape and "             "out_backprop to match, got ",             tensor_in_shape.dim_size(4), " and ", out_backprop.dim_size(4)));      output->flat<T>().setZero();     std::array<int64, 3> input_size = {{tensor_in_shape.dim_size(3),                                         tensor_in_shape.dim_size(2),

Patch:
patch: @@ -383,6 +383,19 @@ struct LaunchAvgPooling3dGradOp<CPUDevice, T> {
                      const std::array<int64, 3>& output_shape,
                      const std::array<int64, 3>& padding,
                      TensorFormat data_format, Tensor* output) {
+    OP_REQUIRES(
+        context, tensor_in_shape.dim_size(0) == out_backprop.dim_size(0),
+        errors::InvalidArgument(
+            "Expected first dimension of tensor_in_shape and "
+            "out_backprop to match, got ",
+            tensor_in_shape.dim_size(0), " and ", out_backprop.dim_size(0)));
+    OP_REQUIRES(
+        context, tensor_in_shape.dim_size(4) == out_backprop.dim_size(4),
+        errors::InvalidArgument(
+            "Expected last dimension of tensor_in_shape and "
+            "out_backprop to match, got ",
+            tensor_in_shape.dim_size(4), " and ", out_backprop.dim_size(4)));
+
     output->flat<T>().setZero();
     std::array<int64, 3> input_size = {{tensor_in_shape.dim_size(3),
                                         tensor_in_shape.dim_size(2),
path: tensorflow/core/kernels/pooling_ops_3d.cc