cummax_cummin_helper Class — pytorch Architecture

Source Code

aten/src/ATen/native/ReduceOps.cpp lines 790–805

template<typename T1, typename T2, typename Operation>
static void cummax_cummin_helper(const T1* self_data, T1* values_data, T2* indices_data,
          int self_dim_size, int self_stride, int values_stride, int indices_stride) {
      Operation op;
      T1 out = c10::load(self_data);
      int idx = 0;
      for (const auto i : c10::irange(self_dim_size)) {
        T1 curr_elem = c10::load(&self_data[i*self_stride]);
        if(isnan_(curr_elem) || (!isnan_(out) && op(curr_elem, out))) {
            out = curr_elem;
            idx = i;
        }
        values_data[i*values_stride] = out;
        indices_data[i*indices_stride] = idx;
      }
}