CVE-2021-41218 (published 2021-11-05T22:15Z)

TensorFlow is an open source platform for machine learning. In affected versions the shape inference code for `AllToAll` can be made to execute a division by 0. This occurs whenever the `split_count` argument is 0. The fix will be included in TensorFlow 2.7.0. We will also cherrypick this commit on TensorFlow 2.6.1, TensorFlow 2.5.2, and TensorFlow 2.4.4, as these are also affected and still in supported range.

Severity: LOW
CVSS v2 base score: 2.1 (AV:L/AC:L/Au:N/C:N/I:N/A:P)
CVSS v3 base score: 5.5, MEDIUM (CVSS:3.1/AV:L/AC:L/PR:L/UI:N/S:U/C:N/I:N/A:H)

Weakness information:
cwe_id: CWE-369
cwe_name: Divide By Zero
description: The product divides a value by zero.
extended_description: This weakness typically occurs when an unexpected value is provided to the product, or if an error occurs that is not properly detected. It frequently occurs in calculations involving physical dimensions such as size, length, width, and height.
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Bruce Fontaine
author_date: 2021-10-01 12:03:47-07:00
msg: Update TPU AllToAll op to avoid divide by 0.

PiperOrigin-RevId: 400259638
Change-Id: Ic4cfe4fe7159da38caed8044ee005f898e42cd86
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: tpu_cross_replica_ops.cc
change_type: ModificationType.MODIFY
diff: @@ -32,6 +32,7 @@ REGISTER_OP("AllToAll")
     .Attr("split_count: int")
     .SetShapeFn([](InferenceContext* c) {
       ShapeHandle input = c->input(0);
+      ShapeHandle group_assignment = c->input(1);
       if (!c->RankKnown(input)) {
         c->set_output(0, c->UnknownShape());
         return Status::OK();
@@ -42,6 +43,21 @@ REGISTER_OP("AllToAll")
       int split_dimension;
       int split_count;
       TF_RETURN_IF_ERROR(c->GetAttr("split_count", &split_count));
+      if (split_count < 1) {
+        return errors::InvalidArgument("split_count ", split_count,
+                                       " must at least be one.");
+      }
+      if (c->RankKnown(group_assignment) && c->Rank(group_assignment) != 2) {
+        return errors::InvalidArgument("group_assignment must have rank 2.");
+      }
+      DimensionHandle num_replicas_per_group = c->Dim(group_assignment, 1);
+      if (c->ValueKnown(num_replicas_per_group) &&
+          (c->Value(num_replicas_per_group) != split_count)) {
+        return errors::InvalidArgument(
+            "split_count ", split_count,
+            " must equal the size of the second dimension of group_assignment ",
+            c->Value(num_replicas_per_group));
+      }
 
       TF_RETURN_IF_ERROR(c->GetAttr("concat_dimension", &concat_dimension));
 
@@ -65,6 +81,12 @@ REGISTER_OP("AllToAll")
           dims[i] = c->MakeDim(c->Value(dims[i]) * split_count);
         }
         if (i == split_dimension) {
+          if (c->ValueKnown(dims[i]) &&
+              (c->Value(dims[i]) % split_count != 0)) {
+            return errors::InvalidArgument(
+                "input dimension ", c->Value(dims[i]),
+                " not divisible by split_count ", split_count);
+          }
           dims[i] = c->MakeDim(c->Value(dims[i]) / split_count);
         }
       }

programming_language: C++
filename: tpu_test.py
change_type: ModificationType.MODIFY
diff: @@ -32,6 +32,7 @@ from tensorflow.python.platform import test
 from tensorflow.python.tpu import tpu
 from tensorflow.python.tpu import tpu_feed
 from tensorflow.python.tpu import training_loop
+from tensorflow.python.tpu.ops import tpu_ops
 
 
 class TPUContextTest(test.TestCase):
@@ -165,6 +166,51 @@ class TPUGraphPruneTest(test.TestCase):
         graph.get_operation_by_name("import/y").get_attr(
             tpu._TPU_REPLICATE_ATTR)
 
+
+class TPUOpsTest(test.TestCase):
+
+  def test_all_to_all_zero_split_count(self):
+    with self.assertRaisesRegex(
+        ValueError, "split_count 0 must at least be one"):
+      tpu_ops.all_to_all(
+          x=[0.0, 0.1652, 0.6543],
+          group_assignment=[1, -1],
+          concat_dimension=0,
+          split_dimension=0,
+          split_count=0)
+
+  def test_all_to_all_group_assignment_wrong_shape(self):
+    with self.assertRaisesRegex(
+        ValueError, "group_assignment must have rank 2"):
+      tpu_ops.all_to_all(
+          x=[0.0, 0.1652, 0.6543],
+          group_assignment=[1, -1],
+          concat_dimension=0,
+          split_dimension=0,
+          split_count=2)
+
+  def test_all_to_all_split_count_not_equal_to_group_assignment_shape(self):
+    with self.assertRaisesRegex(
+        ValueError, "split_count 1 must equal the size of the second dimension "
+        "of group_assignment 2"):
+      tpu_ops.all_to_all(
+          x=[0.0, 0.1652, 0.6543],
+          group_assignment=[[0, 1], [2, 3]],
+          concat_dimension=0,
+          split_dimension=0,
+          split_count=1)
+
+  def test_all_to_all_split_count_not_divide_input_shape(self):
+    with self.assertRaisesRegex(
+        ValueError, "input dimension 3 not divisible by split_count 2"):
+      tpu_ops.all_to_all(
+          x=[[0.0], [0.1652], [0.6543]],
+          group_assignment=[[0, 1], [2, 3]],
+          concat_dimension=1,
+          split_dimension=0,
+          split_count=2)
+
+
 def do_einsum():
   a = array_ops.placeholder(dtype=dtypes.float32, name="a", shape=[2, 3, 4])
   b = array_ops.placeholder(dtype=dtypes.float32, name="b", shape=[2, 4, 5])

programming_language: Python