message: Use the correct byte size of the operand to the comparator function.
  So far, we always used the byte size of the first operand type for all operand
  types. This breaks if operand types have different number of bytes, e.g. s32
  and pred.
diff: @@ -70,11 +70,11 @@ TF_ATTRIBUTE_NO_SANITIZE_MEMORY void __xla_cpu_runtime_KeyValueSort(
  index % sort_dimension_offset +
  (index - index % sort_dimension_offset) * sort_dimension_elements;
  auto compare_function = [&](int64 a, int64 b) -> bool {
  + for (int32 i = 0; i < values_count; ++i) {
  int64 memory_index_lhs = (base_offset + a * sort_dimension_offset) *
  - values_primitive_type_size_in_bytes[0];
  + values_primitive_type_size_in_bytes[i];
  int64 memory_index_rhs = (base_offset + b * sort_dimension_offset) *
  - values_primitive_type_size_in_bytes[0];
  - for (int32 i = 0; i < values_count; ++i) {
  + values_primitive_type_size_in_bytes[i];
  comparison_values[i * 2] = values[i] + memory_index_lhs;
  comparison_values[i * 2 + 1] = values[i] + memory_index_rhs;
  }
  