message: Support int64 values for num_segments in UnsortedSegmentReductionOp.
diff: @@ -478,8 +478,9 @@ class UnsortedSegmentReductionOp : public OpKernel {
  return;
  }
  const auto segment_flat = segment_ids.flat<Index>();
  - const Index output_rows =
  - internal::SubtleMustCopy(num_segments.scalar<int32>()());
  + const Index output_rows = internal::SubtleMustCopy(static_cast<Index>(
  + num_segments.dtype() == DT_INT32 ? num_segments.scalar<int32>()()
  + : num_segments.scalar<int64>()()));
  OP_REQUIRES(context, output_rows >= 0,
  errors::InvalidArgument("Input num_segments == ", output_rows,
  " must not be negative."));
  