Fix XLA DenseBincount negative input validation for runtime tensors - #122706

Code Review

This pull request adds support for handling runtime negative inputs in the XLA dense_bincount operator by mapping negative indices to an out-of-bounds value (output_size) so they are silently ignored by XLA Scatter, along with adding corresponding unit tests. The review feedback points out that also zeroing out the updates for these negative values is redundant and introduces unnecessary performance overhead, suggesting the removal of these redundant operations and updating the comments accordingly.

Correctness & Safety

The approach of using output_size (Out of Bounds - OOB) instead of clamping to 0 is excellent. It correctly avoids overwriting valid data (e.g., bin 0) in cases where binary_output=True (Scatter Assign semantics). Since XLA Scatter drops OOB updates, setting negative indices to a positive OOB value like output_size is a standard and robust way to ignore them.

Redundancy in Updates Zeroing

The PR description notes step 3 (zeroing updates) as a "belt-and-suspenders safety measure".

• If Scatter strictly adheres to XLA semantics and drops OOB indices (which includes output_size), zeroing updates is redundant because the update value is never applied.
• Zeroing updates introduces minor dynamic overhead (Mul, ConvertElementType).
• Caution: If there were a backend bug where OOB indices were accidentally clamped or routed to a valid bin (e.g., bin 0), zeroing out updates could actually make the bug worse for ScatterAssign (Assign 0 overwriting valid data) than leaving it as-is (Assign 1 or weight).
• Recommendation: Consider removing the updates zeroing step if performance is critical or if you are confident in standard OOB dropping behavior. The explicit conversion of negative indices to positive OOB (output_size) is sufficient and safer than leaving them negative (which might cause UB or buffer overflow on non-compliant backends).

Consistent Positive OOB

Using output_size is preferred over relying on negative indices being treated as OOB, as some backends or low-level arithmetic might handle signed negative indices inconsistently or interpret them as huge positive indices wrapping around.

Unit Tests

To verify this behavior, the following tests can be added to standard suite (e.g., bincount_op_test.py). These tests verify that runtime negative inputs are ignored gracefully.

  def testNegativeInputRuntimeIgnored(self):
    with self.session():
      with self.test_scope():

        @def_function.function(jit_compile=True)
        def f(arr):
          return gen_math_ops.dense_bincount(
              input=arr, size=3, weights=[]
          )

        # Pass as argument to make it a runtime tensor
        arr = array_ops.constant([0, -1, 2])
        result = f(arr)
        # Expected: 0 -> bin 0, 2 -> bin 2. -1 ignored.
        self.assertAllEqual([1, 0, 1], result)

  def testNegativeInputRuntimeIgnoredBinary(self):
    with self.session():
      with self.test_scope():

        @def_function.function(jit_compile=True)
        def f(arr):
          return gen_math_ops.dense_bincount(
              input=arr, size=3, weights=[], binary_output=True
          )

        arr = array_ops.constant([0, -1, 2])
        result = f(arr)
        self.assertAllEqual([1, 0, 1], result)

  def testNegativeInputRuntime2d(self):
    with self.session():
      with self.test_scope():

        @def_function.function(jit_compile=True)
        def f(arr):
          return gen_math_ops.dense_bincount(
              input=arr, size=3, weights=[]
          )

        # 2D input: [[0, -1, 2], [1, 2, -1]]
        arr = array_ops.constant([[0, -1, 2], [1, 2, -1]])
        result = f(arr)
        # Expected output shape: [2, 3]
        # Row 0: 0 -> bin 0, 2 -> bin 2. -1 ignored. -> [1, 0, 1]
        # Row 1: 1 -> bin 1, 2 -> bin 2. -1 ignored. -> [0, 1, 1]
        self.assertAllEqual([[1, 0, 1], [0, 1, 1]], result)

We have tested the changes internally and they look good.

Verification Results

• CPU Tests: Passed.
• GPU Tests: Passed (verified with CUDA backend).
• A100 Tests: Passed.

Logic Review

The approach of mapping negative indices in runtime tensors to output_size (out of bounds) using xla::Select successfully prevents corruption and ensures that negative values are ignored, matching the behavior of the non-XLA kernels.

Removing the redundant updates zeroing for negative indices (in your 4th commit) is also correct as XLA Scatter handles out-of-bounds updates silently.

Code Quality & Style

Please ensure your code follows standard formatting rules. Running clang-format on your C++ changes (bincount_op.cc) will help maintain consistency.

Implementation Details

Your approach of masking negative indices to output_size (out of bounds) so that XLA Scatter silently ignores them is valid and works correctly for both 1D and 2D inputs.

However, running these tests exposed a couple of issues that were addressed:

1. DenseBincount Shape Inference: The shape function in tensorflow/core/ops/math_ops.cc was missing a fallback for unknown rank inputs. This caused a crash (Check failed: output(i).IsSet()) during tracing of your new tests. Added a fallback to set output to UnknownShape() for unknown rank inputs.

   // In tensorflow/core/ops/math_ops.cc under REGISTER_OP("DenseBincount")
   // ...
         if (c->Rank(c->input(0)) == 1 || c->Rank(c->input(0)) == 0) {
           c->set_output(0, c->MakeShape({size_val}));
         } else if (c->Rank(c->input(0)) == 2) {
           c->set_output(0, c->MakeShape({c->Dim(c->input(0), 0), size_val}));
   +     } else {
   +       c->set_output(0, c->UnknownShape());
         }
         return absl::OkStatus();
       });

2. Test Wrappers: In bincount_op_test.py, the new tests required with self.session(): wrappers to avoid test_session not supported errors under XLATestCase.

   # In tensorflow/compiler/tests/bincount_op_test.py
   # ...
     def testNegativeInputRuntimeIgnored(self):
   -    with self.test_scope():
   +    with self.session():
   +      with self.test_scope():
           @def_function.function(
               jit_compile=True,
               input_signature=[tensor_spec.TensorSpec([5], dtypes.int32)])
           def f(arr):
             return gen_math_ops.dense_bincount(
                 input=arr, size=5, weights=[]
             )

   (Please apply this with self.session(): wrapper to all three new tests: testNegativeInputRuntimeIgnored, testNegativeInputRuntimeIgnoredBinary, and testNegativeInputRuntime2d)

Please include these fixes in your GitHub PR to ensure tests pass successfully.