Merge pull request #63344 from tensorflow:Surya_MatrixDiagV3_checkfai…

…l_fix PiperOrigin-RevId: 625362770
tensorflow · Apr 16, 2024 · adb4c10 · adb4c10
2 parents 59bf209 + 7186b5b
commit adb4c10
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/tensorflow/core/kernels/linalg/matrix_diag_op.cc b/tensorflow/core/kernels/linalg/matrix_diag_op.cc
@@ -236,12 +236,13 @@ class MatrixDiagOp : public OpKernel {
         errors::InvalidArgument(
             "lower_diag_index must not be larger than upper_diag_index: ",
             lower_diag_index, " > ", upper_diag_index));
-    OP_REQUIRES(context,
-                lower_diag_index == upper_diag_index ||
-                    diagonal_shape.dim_size(diag_rank - 2) == num_diags,
-                errors::InvalidArgument(
-                    "The number of diagonals provided in the input does not "
-                    "match the lower_diag_index and upper_diag_index range."));
+    OP_REQUIRES(
+        context,
+        lower_diag_index == upper_diag_index ||
+            diagonal_shape.dim_size(std::max(diag_rank - 2, 0)) == num_diags,
+        errors::InvalidArgument(
+            "The number of diagonals provided in the input does not "
+            "match the lower_diag_index and upper_diag_index range."));
 
     const Eigen::Index max_diag_len = diagonal_shape.dim_size(diag_rank - 1);
     const Eigen::Index min_num_rows =