Merge pull request #129180 from destro4evr-rgb:fix/ctc-loss-gpu-oob-batch-index

PiperOrigin-RevId: 1003702391

tensorflow/core/kernels/ctc_loss_op.cc

@@ -54,17 +54,6 @@ using se::StreamExecutor;
 using se::dnn::RnnStateTensorDescriptor;
 using se::dnn::ToDataType;
 
-template <typename T>
-void DoHistogram(OpKernelContext* ctx, const Tensor* labels_indices,
-                 int num_indices, int batch_size,
-                 std::vector<int>* labels_lengths) {
-  const T* h_in = labels_indices->flat<T>().data();
-  for (int i = 0; i < num_indices; i++) {
-    const T& key = h_in[i * 2];
-    (*labels_lengths)[key]++;
-  }
-}
-
 }  // end namespace
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
@@ -283,6 +272,10 @@ class CTCLossOpGPU : public OpKernel {
                 absl::InvalidArgumentError("labels_indices is not a matrix"));
     OP_REQUIRES(ctx, TensorShapeUtils::IsVector(labels_values->shape()),
                 absl::InvalidArgumentError("labels_values is not a vector"));
+    OP_REQUIRES(ctx, labels_indices->dim_size(1) == 2,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "labels_indices second dimension must be 2. Received ",
+                    labels_indices->dim_size(1))));
 
     const TensorShape& inputs_shape = inputs->shape();
     const int64_t max_time_raw = inputs_shape.dim_size(0);
@@ -318,6 +311,18 @@ class CTCLossOpGPU : public OpKernel {
     OP_REQUIRES(ctx, batch_size != 0,
                 absl::InvalidArgumentError("batch_size must not be 0"));
 
+    // Validate batch indices and convert labels_indices to labels_lengths.
+    std::vector<int> labels_lengths(batch_size, 0);
+    const auto labels_mat = labels_indices->matrix<int64_t>();
+    for (int64_t i = 0; i < num_indices; ++i) {
+      const int64_t batch_idx = labels_mat(i, 0);
+      OP_REQUIRES(ctx, FastBoundsCheck(batch_idx, batch_size),
+                  absl::InvalidArgumentError(
+                      absl::StrCat("labels_indices batch index ", batch_idx,
+                                   " out of range [0, ", batch_size, ")")));
+      ++labels_lengths[batch_idx];
+    }
+
     Tensor* loss = nullptr;
     OP_REQUIRES_OK(ctx, ctx->allocate_output("loss", seq_len->shape(), &loss));
 
@@ -325,11 +330,6 @@ class CTCLossOpGPU : public OpKernel {
     OP_REQUIRES_OK(ctx,
                    ctx->allocate_output("gradient", inputs_shape, &gradient));
 
-    // Convert the labels_indices to labels_lengths.
-    std::vector<int> labels_lengths(batch_size, 0);
-    DoHistogram<int64_t>(ctx, labels_indices, num_indices, batch_size,
-                         &labels_lengths);
-
     StreamExecutor* executor = ctx->op_device_context()->stream()->parent();
     se::dnn::DataType data_type = ToDataType<float>::value;
 

tensorflow/python/kernel_tests/nn_ops/ctc_loss_op_test.py

@@ -1236,6 +1236,69 @@ class CTCLossTestV3(test.TestCase, parameterized.TestCase):
     self.assertAllClose(loss, ref_loss, atol=1e-6)
     self.assertAllClose(grad, ref_grad, atol=2e-6)
 
+  @test_util.run_v2_only
+  def testGPULabelsIndicesOOBBatchIndexRaisesError(self):
+    """GPU CTCLossV2 rejects out-of-range labels_indices batch indices."""
+    if not test.is_gpu_available():
+      self.skipTest("Need GPU for testing.")
+    batch_size = 2
+    num_frames = 4
+    num_classes = 3
+    labels_values = constant_op.constant([1, 1], dtype=dtypes.int32)
+    logits = constant_op.constant(
+        np.zeros([num_frames, batch_size, num_classes], dtype=np.float32)
+    )
+    sequence_length = constant_op.constant(
+        [num_frames] * batch_size, dtype=dtypes.int32
+    )
+    for bad_batch in (99, -1):
+      labels_indices = constant_op.constant(
+          [[bad_batch, 0], [0, 1]], dtype=dtypes.int64
+      )
+      with ops.device("/GPU:0"):
+        with self.assertRaisesRegex(
+            errors_impl.InvalidArgumentError,
+            rf"labels_indices batch index {bad_batch} out of range",
+        ):
+          self.evaluate(
+              gen_ctc_ops.ctc_loss_v2(
+                  logits,
+                  labels_indices,
+                  labels_values,
+                  sequence_length,
+                  preprocess_collapse_repeated=False,
+                  ctc_merge_repeated=True,
+                  ignore_longer_outputs_than_inputs=False,
+              )
+          )
+
+  @test_util.run_v2_only
+  def testGPULabelsIndicesWrongColumnCountRaisesError(self):
+    """GPU CTCLossV2 requires labels_indices to have 2 columns."""
+    if not test.is_gpu_available():
+      self.skipTest("Need GPU for testing.")
+    logits = constant_op.constant(np.zeros([4, 2, 3], dtype=np.float32))
+    labels_values = constant_op.constant([1, 1], dtype=dtypes.int32)
+    sequence_length = constant_op.constant([4, 4], dtype=dtypes.int32)
+    for bad_indices in ([[0], [1]], [[0, 0, 99], [1, 0, 0]]):
+      labels_indices = constant_op.constant(bad_indices, dtype=dtypes.int64)
+      with ops.device("/GPU:0"):
+        with self.assertRaisesRegex(
+            errors_impl.InvalidArgumentError,
+            r"labels_indices second dimension must be 2",
+        ):
+          self.evaluate(
+              gen_ctc_ops.ctc_loss_v2(
+                  logits,
+                  labels_indices,
+                  labels_values,
+                  sequence_length,
+                  preprocess_collapse_repeated=False,
+                  ctc_merge_repeated=True,
+                  ignore_longer_outputs_than_inputs=False,
+              )
+          )
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class CTCLossDeterministicTest(test.TestCase, parameterized.TestCase):