CVE-2021-41218: Divide By Zero (CWE-369)
value: 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.
CWE description: The product divides a value by zero.
CVSS v3 base score: 5.5
Published: 2021-11-05T22:15Z
Repository: https://github.com/tensorflow/tensorflow
Language: Python
Files changed: tensorflow/core/ops/tpu_cross_replica_ops.cc, tensorflow/python/tpu/tpu_test.py
Commit date: 2021-10-01 12:15:18 -0700
Commit message: Update TPU AllToAll op to avoid divide by 0.

PiperOrigin-RevId: 400259638
Change-Id: Ic4cfe4fe7159da38caed8044ee005f898e42cd86

Diff: diff --git a/tensorflow/core/ops/tpu_cross_replica_ops.cc b/tensorflow/core/ops/tpu_cross_replica_ops.cc
index 045b2dffaf615..74be53c921a98 100644
--- a/tensorflow/core/ops/tpu_cross_replica_ops.cc
+++ b/tensorflow/core/ops/tpu_cross_replica_ops.cc
@@ -30,20 +30,36 @@ REGISTER_OP("AllToAll")
     .Attr("concat_dimension: int")
     .Attr("split_dimension: int")
     .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();
       }
 
       int64_t rank = c->Rank(input);
       int concat_dimension;
       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));
 
       if (concat_dimension < 0 || concat_dimension >= rank) {
         return errors::InvalidArgument("concat_dimension ", concat_dimension,
@@ -63,10 +79,16 @@ REGISTER_OP("AllToAll")
         dims[i] = c->Dim(input, i);
         if (i == concat_dimension) {
           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);
         }
       }
 
       c->set_output(0, c->MakeShape(dims));
diff --git a/tensorflow/python/tpu/tpu_test.py b/tensorflow/python/tpu/tpu_test.py
index 10707fc3d5145..48139171e1abd 100644
--- a/tensorflow/python/tpu/tpu_test.py
+++ b/tensorflow/python/tpu/tpu_test.py
@@ -30,10 +30,11 @@ from tensorflow.python.ops import special_math_ops
 from tensorflow.python.ops import variable_scope
 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):
 
   def testIsInContext(self):
@@ -163,10 +164,55 @@ class TPUGraphPruneTest(test.TestCase):
           ValueError,
           "Operation \'import/y\' has no attr named \'_tpu_replicate\'"):
         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])
   return special_math_ops.einsum("abc,acd->abd", a, b)
 

Fixed code:       ShapeHandle group_assignment = c->input(1);
      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));
      }
          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);
          }
from tensorflow.python.tpu.ops import tpu_ops

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)