Vulnerability: Heap OOB read in `tf.raw_ops.Dequantize` ### Impact
Due to lack of validation in `tf.raw_ops.Dequantize`, an attacker can trigger a read from outside of bounds of heap allocated data:

```python
import tensorflow as tf

input_tensor=tf.constant(
  [75, 75, 75, 75, -6, -9, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10,\
  -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10,\
  -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10, -10,\
  -10, -10, -10, -10], shape=[5, 10], dtype=tf.int32)
input_tensor=tf.cast(input_tensor, dtype=tf.quint8)
min_range = tf.constant([-10], shape=[1], dtype=tf.float32)
max_range = tf.constant([24, 758, 758, 758, 758], shape=[5], dtype=tf.float32)
  
tf.raw_ops.Dequantize( 
  input=input_tensor, min_range=min_range, max_range=max_range, mode='SCALED',
  narrow_range=True, axis=0, dtype=tf.dtypes.float32)
```

The [implementation](https://github.com/tensorflow/tensorflow/blob/26003593aa94b1742f34dc22ce88a1e17776a67d/tensorflow/core/kernels/dequantize_op.cc#L106-L131) accesses the `min_range` and `max_range` tensors in parallel but fails to check that they have the same shape:

```cc
if (num_slices == 1) {
  const float min_range = input_min_tensor.flat<float>()(0);
  const float max_range = input_max_tensor.flat<float>()(0);
  DequantizeTensor(ctx, input, min_range, max_range, &float_output);
} else {
  ...
  auto min_ranges = input_min_tensor.vec<float>();
  auto max_ranges = input_max_tensor.vec<float>();
  for (int i = 0; i < num_slices; ++i) {
    DequantizeSlice(ctx->eigen_device<Device>(), ctx,
                    input_tensor.template chip<1>(i), min_ranges(i),
                    max_ranges(i), output_tensor.template chip<1>(i));
    ...
  }
}
```

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

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 Yakun Zhang and Ying Wang of Baidu X-Team.
Vulnerability ID: GHSA-c45w-2wxr-pp53
CWE: CWE-125
Severity score: 2.5
Project: https://github.com/tensorflow/tensorflow
Published: 2021-05-21
File: tensorflow/core/kernels/dequantize_op.cc
Commit message: Fix heap OOB read in dequantize op.

Also fixes SEGV in same op

PiperOrigin-RevId: 372437896
Change-Id: I135e94d360c2a1ce374c10f7e0fed1af603dbc02
Commit author: Mihai Maruseac

Vulnerable version (before fix):
 class DequantizeOp : public OpKernel {     if (axis_ > -1) {       num_slices = input.dim_size(axis_);     }      Tensor* output = nullptr;     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));

Fixed version (after fix):
 class DequantizeOp : public OpKernel {     if (axis_ > -1) {       num_slices = input.dim_size(axis_);     }     OP_REQUIRES(ctx, input_min_tensor.NumElements() == num_slices,                 errors::InvalidArgument(                     "input_min_tensor must have as many elements as input on "                     "the dequantization axis (",                     axis_, "), got ", input_min_tensor.NumElements(),                     ", expected ", num_slices));     OP_REQUIRES(ctx, input_max_tensor.NumElements() == num_slices,                 errors::InvalidArgument(                     "input_max_tensor must have as many elements as input on "                     "the dequantization axis (",                     axis_, "), got ", input_max_tensor.NumElements(),                     ", expected ", num_slices));      Tensor* output = nullptr;     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));

Patch:
patch: @@ -98,6 +98,18 @@ class DequantizeOp : public OpKernel {
     if (axis_ > -1) {
       num_slices = input.dim_size(axis_);
     }
+    OP_REQUIRES(ctx, input_min_tensor.NumElements() == num_slices,
+                errors::InvalidArgument(
+                    "input_min_tensor must have as many elements as input on "
+                    "the dequantization axis (",
+                    axis_, "), got ", input_min_tensor.NumElements(),
+                    ", expected ", num_slices));
+    OP_REQUIRES(ctx, input_max_tensor.NumElements() == num_slices,
+                errors::InvalidArgument(
+                    "input_max_tensor must have as many elements as input on "
+                    "the dequantization axis (",
+                    axis_, "), got ", input_max_tensor.NumElements(),
+                    ", expected ", num_slices));
 
     Tensor* output = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));
path: tensorflow/core/kernels/dequantize_op.cc