CVE-2022-36027 (published 2022-09-16T23:15Z)

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.

CVSS v3 base score: 7.5, HIGH (CVSS:3.1/AV:N/AC:L/PR:N/UI:N/S:U/C:N/I:N/A:H)

Weakness information:
cwe_id: NVD-CWE-noinfo
cwe_name: Insufficient Information
description: There is insufficient information about the issue to classify it; details are unkown or unspecified.
extended_description: Insufficient Information

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Jaesung Chung
author_date: 2022-07-25 07:28:07-07:00
msg: 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
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: quantization_utils.cc
change_type: ModificationType.MODIFY
diff: @@ -168,12 +168,10 @@ quant::UniformQuantizedPerAxisType ResetAxisAndBroadcast(
         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 {
@@ -343,6 +341,10 @@ TypeAttr CastQuantizedTypeAttrFromExpressedType(Builder builder,
   // 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);
   }

programming_language: C++
filename: prepare_tf.cc
change_type: ModificationType.MODIFY
diff: @@ -1370,6 +1370,8 @@ void PrepareTFPass::runOnOperation() {
 
   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
@@ -1399,6 +1401,8 @@ void PrepareTFPass::runOnOperation() {
            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));
 }

programming_language: C++
filename: lite_v2_test.py
change_type: ModificationType.MODIFY
diff: @@ -2311,6 +2311,44 @@ class FromSavedModelTest(lite_v2_test_util.ModelTest):
         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')

programming_language: Python