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

PiperOrigin-RevId: 400259638
Change-Id: Ic4cfe4fe7159da38caed8044ee005f898e42cd86


Vulnerable code:

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)