Fix heap buffer overflow in UnsortedSegmentSum.

When Index=int32, data_size and num_segments were truncated from int64 to int32. This truncation can produce negative numbers, which causes UnsortedSegmentFunctor to access out of bounds memory.

Also:
- Switches some indexing calculations to int64 to avoid signed integer overflow when either the input or output tensors have more than 2**31 - 1 elements.
- Fixes a range check error in the GPU kernel. The segment ID was checked against an upper bound measured in elements, not segments.
PiperOrigin-RevId: 256451663