From 8a793b5d7f59e37ac7f3cd0954a750a2fe76bad4 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Thu, 29 Jul 2021 18:23:45 -0700
Subject: [PATCH] Prevent division by 0 in common shape functions.

PiperOrigin-RevId: 387712197
Change-Id: Id25c7460e35b68aeeeac23b9a88e455b443ee149
---
 tensorflow/core/framework/common_shape_fns.cc | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tensorflow/core/framework/common_shape_fns.cc b/tensorflow/core/framework/common_shape_fns.cc
index b9efddf4..e25d5581 100644
--- a/tensorflow/core/framework/common_shape_fns.cc
+++ b/tensorflow/core/framework/common_shape_fns.cc
@@ -659,6 +659,8 @@ Status Conv2DShapeImpl(shape_inference::InferenceContext* c,
   if (c->ValueKnown(input_depth_dim) && c->ValueKnown(filter_input_depth_dim)) {
     int64 input_depth_value = c->Value(input_depth_dim),
           filter_input_depth_value = c->Value(filter_input_depth_dim);
+    if (filter_input_depth_value == 0)
+      return errors::InvalidArgument("Depth of filter must not be 0");
     if (input_depth_value % filter_input_depth_value != 0)
       return errors::InvalidArgument(
           "Depth of input (", input_depth_value,
@@ -668,6 +670,8 @@ Status Conv2DShapeImpl(shape_inference::InferenceContext* c,
       int64 num_groups = input_depth_value / filter_input_depth_value;
       if (c->ValueKnown(output_depth_dim)) {
         int64 output_depth_value = c->Value(output_depth_dim);
+	if (num_groups == 0)
+	  return errors::InvalidArgument("Number of groups must not be 0");
         if (output_depth_value % num_groups != 0)
           return errors::InvalidArgument(
               "Depth of output (", output_depth_value,
@@ -798,6 +802,8 @@ Status Conv3DShape(shape_inference::InferenceContext* c) {
   if (c->ValueKnown(input_depth_dim) && c->ValueKnown(filter_input_depth_dim)) {
     int64 input_depth_value = c->Value(input_depth_dim),
           filter_input_depth_value = c->Value(filter_input_depth_dim);
+    if (filter_input_depth_value == 0)
+      return errors::InvalidArgument("Depth of filter must not be 0");
     if (input_depth_value % filter_input_depth_value != 0)
       return errors::InvalidArgument(
           "Depth of input (", input_depth_value,
@@ -807,6 +813,8 @@ Status Conv3DShape(shape_inference::InferenceContext* c) {
       int64 num_groups = input_depth_value / filter_input_depth_value;
       if (c->ValueKnown(output_depth_dim)) {
         int64 output_depth_value = c->Value(output_depth_dim);
+	if (num_groups == 0)
+          return errors::InvalidArgument("Number of groups must not be 0");
         if (output_depth_value % num_groups != 0)
           return errors::InvalidArgument(
               "Depth of output (", output_depth_value,
@@ -2364,6 +2372,9 @@ Status SparseReduceShapeFn(InferenceContext* c) {
 
     int64 ndims = shape_vec.size();
     absl::flat_hash_set<int64> axes;
+    if (ndims == 0)
+      return errors::InvalidArgument(
+	  "Number of dims in shape tensor must not be 0");
     for (int i = 0; i < axes_vec.size(); i++) {
       axes.insert((axes_vec(i) + ndims) % ndims);
     }
-- 
2.27.0