CVE-2022-36027: Insufficient Information (NVD-CWE-noinfo)
value: TensorFlow is an open source platform for machine learning. When converting transposed convolutions using per-channel weight quantization the converter segfaults and crashes the Python process. We have patched the issue in GitHub commit aa0b852a4588cea4d36b74feb05d93055540b450. The fix will be included in TensorFlow 2.10.0. We will also cherrypick this commit on TensorFlow 2.9.1, TensorFlow 2.8.1, and TensorFlow 2.7.2, as these are also affected and still in supported range. There are no known workarounds for this issue.
CWE description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
CVSS v3 base score: 7.5
Published: 2022-09-16T23:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: C++
Files changed: tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc, tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc, tensorflow/lite/python/lite_v2_test.py
Commit date: 2022-07-25 07:32:03 -0700
Commit message: Skip reordering dq-q patterns when the new quantization dimension is unknown.

It also adds the canonicalization patterns to remove redundant TF reshape ops
to infer the quantization dimension from the user ops (e.g, conv2d, depthwise
conv2d) over those redundant TF reshape ops.

PiperOrigin-RevId: 463080509

Diff: diff --git a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
index 0bee61e913fbe..d9a7534a21422 100644
--- a/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/quantization_utils.cc
@@ -166,16 +166,14 @@ quant::UniformQuantizedPerAxisType ResetAxisAndBroadcast(
     // have to repeat each elements in the `scales` locally dim[3] times.
     if (BroadcastVector<double>(shaped.getDimSize(quant_dim), scales) ||
         BroadcastVector<int64_t>(shaped.getDimSize(quant_dim), zero_points)) {
       return {};
     }
-  } else if ((new_shape.size() == shape.size() + 1) && new_shape.back() == 1) {
-    // This is a trivial shift left, then we shift the quant_dim as well.
-    if (std::equal(shape.begin(), shape.end(), new_shape.begin()) &&
-        quant_dim == -1) {
-      quant_dim = shape.size() + quant_dim;
-    } else {
+  } else if ((new_shape.size() == shape.size() + 1) && new_shape.front() == 1) {
+    // Handle the [A, B, C] -> [1, A, B, C] reshape case.
+    if (!(std::equal(shape.begin(), shape.end(), new_shape.begin() + 1) &&
+          quant_dim == new_shape.size() - 1)) {
       return {};
     }
   } else {
     return {};
   }
@@ -341,10 +339,14 @@ TypeAttr CastQuantizedTypeAttrFromExpressedType(Builder builder,
   auto qtype = src_ele_type.dyn_cast<quant::QuantizedType>();
 
   // Reset the quantization dimensions if it is per-axis.
   if (auto per_axis =
           qtype.dyn_cast_or_null<quant::UniformQuantizedPerAxisType>()) {
+    // For the pass-through ops, we don't know which the dimension will be the
+    // new quantization dimension. Only if the new quantization dimension can
+    // be inferred, it is safe to reset the per-axis quantized type.
+    if (axis == -1) return {};
     qtype =
         ResetAxisAndBroadcast(source_type.getShape(), per_axis, target, axis);
   }
   if (!qtype) return {};
   Type final_type = qtype.castFromExpressedType(target);
diff --git a/tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc b/tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc
index f1dafc535b2b3..5cf5004e78f7e 100644
--- a/tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/prepare_tf.cc
@@ -1368,10 +1368,12 @@ void PrepareTFPass::runOnOperation() {
   patterns.add<ConvertTFDilatedConvOp<TF::Conv2DOp>, FusedBatchNormV3Pat,
                ConvertTFDilatedConvOp<TF::DepthwiseConv2dNativeOp>>(ctx);
 
   patterns.add<RemoveIdentity>(ctx);
   TFL::populateWithGenerated(patterns);
+  // Remove redundant reshape ops.
+  TF::ReshapeOp::getCanonicalizationPatterns(patterns, ctx);
   // TODO(karimnosseir): Split to separate pass probably after
   // deciding on long term plan for this optimization.
   // This will allow optimizing any TF_Mul->TF_Conv in the graph
   // and any expanded from FusedBatchNorm. We need to do this
   // before converting TF_Conv to TFL_Conv
@@ -1397,10 +1399,12 @@ void PrepareTFPass::runOnOperation() {
   phase_2_patterns
       .add<TF::ConvertTFEinsumOp, ConvertTFBroadcastTo, ConvertTFStridedSlice,
            ConvertRfftToRfft2d, RemoveIdentity>(ctx);
   phase_2_patterns.add<ConvertTFConv2D, ConvertTFDepthwiseConv2dNative>(
       ctx, allow_bf16_and_f16_type_legalization_);
+  // Remove redundant reshape ops.
+  TF::ReshapeOp::getCanonicalizationPatterns(phase_2_patterns, ctx);
 
   (void)applyPatternsAndFoldGreedily(func, std::move(phase_2_patterns));
 }
 
 }  // namespace
diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py
index bcdad1db9ed87..b049832c7fc28 100644
--- a/tensorflow/lite/python/lite_v2_test.py
+++ b/tensorflow/lite/python/lite_v2_test.py
@@ -2309,10 +2309,48 @@ class FromSavedModelTest(lite_v2_test_util.ModelTest):
     self.assertAllClose(expected_value, actual_value, atol=1e-1)
     self.assertEqual(
         list(output_details[0]['shape_signature']),
         list(model.layers[-1].output_shape))
 
+  @test_util.run_v2_only
+  def testKerasConv2DTransposedWithMismatchQuantizedAxes(self):
+
+    class QuantConv2DTransposed(tf.keras.layers.Layer):
+
+      def build(self, input_shape):
+        self.kernel = self.add_weight('kernel', [3, 3, input_shape[-1], 24])
+
+      def call(self, inputs):
+        filters = tf.quantization.fake_quant_with_min_max_vars_per_channel(
+            self.kernel,
+            -3.0 * tf.ones([24]),
+            3.0 * tf.ones([24]),
+            narrow_range=True)
+        filters = tf.transpose(filters, (0, 1, 3, 2))
+        return tf.nn.conv2d_transpose(inputs, filters, [*inputs.shape[:-1], 24],
+                                      1)
+
+    inp = tf.keras.Input(shape=(6, 8, 48), batch_size=1)
+    x = tf.quantization.fake_quant_with_min_max_vars(
+        inp, -3.0, 3.0, narrow_range=True)
+    x = QuantConv2DTransposed()(x)
+    x = tf.quantization.fake_quant_with_min_max_vars(
+        x, -3.0, 3.0, narrow_range=True)
+
+    model = tf.keras.Model(inp, x)
+
+    saved_model_dir = os.path.join(self.get_temp_dir(),
+                                   'keras_conv2d_transpose')
+    model.save(saved_model_dir)
+    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
+    converter.optimizations = [tf.lite.Optimize.DEFAULT]
+
+    with self.assertRaises(convert.ConverterError) as error:
+      _ = converter.convert()
+    self.assertIn('mismatched quantized axes of input and output',
+                  str(error.exception))
+
   def _createModelWithInputShape(self, shape):
     """Create a simple SavedModel with a certain shape."""
     saved_model_dir = os.path.join(self.get_temp_dir(), 'input_shape_model')
     with tf.Graph().as_default():
       with tf.compat.v1.Session() as sess:

Vulnerable code:   } else if ((new_shape.size() == shape.size() + 1) && new_shape.back() == 1) {
    // This is a trivial shift left, then we shift the quant_dim as well.
    if (std::equal(shape.begin(), shape.end(), new_shape.begin()) &&
        quant_dim == -1) {
      quant_dim = shape.size() + quant_dim;
    } else {
Fixed code:   } else if ((new_shape.size() == shape.size() + 1) && new_shape.front() == 1) {
    // Handle the [A, B, C] -> [1, A, B, C] reshape case.
    if (!(std::equal(shape.begin(), shape.end(), new_shape.begin() + 1) &&
          quant_dim == new_shape.size() - 1)) {
    // For the pass-through ops, we don't know which the dimension will be the
    // new quantization dimension. Only if the new quantization dimension can
    // be inferred, it is safe to reset the per-axis quantized type.
    if (axis == -1) return {};
  // Remove redundant reshape ops.
  TF::ReshapeOp::getCanonicalizationPatterns(patterns, ctx);
  // Remove redundant reshape ops.
  TF::ReshapeOp::getCanonicalizationPatterns(phase_2_patterns, ctx);
  @test_util.run_v2_only
  def testKerasConv2DTransposedWithMismatchQuantizedAxes(self):

    class QuantConv2DTransposed(tf.keras.layers.Layer):

      def build(self, input_shape):
        self.kernel = self.add_weight('kernel', [3, 3, input_shape[-1], 24])

      def call(self, inputs):
        filters = tf.quantization.fake_quant_with_min_max_vars_per_channel(
            self.kernel,
            -3.0 * tf.ones([24]),
            3.0 * tf.ones([24]),
            narrow_range=True)
        filters = tf.transpose(filters, (0, 1, 3, 2))
        return tf.nn.conv2d_transpose(inputs, filters, [*inputs.shape[:-1], 24],
                                      1)

    inp = tf.keras.Input(shape=(6, 8, 48), batch_size=1)
    x = tf.quantization.fake_quant_with_min_max_vars(
        inp, -3.0, 3.0, narrow_range=True)
    x = QuantConv2DTransposed()(x)
    x = tf.quantization.fake_quant_with_min_max_vars(
        x, -3.0, 3.0, narrow_range=True)

    model = tf.keras.Model(inp, x)

    saved_model_dir = os.path.join(self.get_temp_dir(),
                                   'keras_conv2d_transpose')
    model.save(saved_model_dir)
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]

    with self.assertRaises(convert.ConverterError) as error:
      _ = converter.convert()
    self.assertIn('mismatched quantized axes of input and output',
                  str(error.exception))