binary_search_strided_rightmost Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseTensorMath.cpp lines 1907–1945

template<typename scalar_t>
static scalar_t binary_search_strided_rightmost(scalar_t search_val, TensorAccessor<scalar_t, 1>& sorted_arr_accessor, int64_t sorted_arr_begin_idx, int64_t length, bool* found) {
  if (length == 0) {
    *found = false;
    return -1;
  }

  int64_t left_ind = 0;
  int64_t right_ind = length - 1;
  // This value should be overwritten in the loop so we use
  // a destructive initial value to ensure disaster if that
  // turns out not to be the case.
  int64_t mid_ind = std::numeric_limits<int64_t>::max();
  bool done_searching = false;

  while (!done_searching) {
    mid_ind = left_ind + (right_ind - left_ind) / 2;
    scalar_t mid_val = sorted_arr_accessor[sorted_arr_begin_idx + mid_ind];

    if (mid_val > search_val) {
      right_ind = mid_ind-1;
    } else if((mid_val == search_val) && (
      (mid_ind == length - 1) || (sorted_arr_accessor[sorted_arr_begin_idx + mid_ind + 1] != search_val)
    )) {
      done_searching = true;
      *found = true;
    } else {
      left_ind = mid_ind+1;
    }

    if (left_ind > right_ind) {
      done_searching = true;
      *found = false;
      mid_ind = -1;
    }
  }

  return mid_ind;
}