compute_counts_uniq Class — pytorch Architecture

Source Code

aten/src/ATen/native/EmbeddingBag.cpp lines 1490–1507

template<typename index_t>
static std::vector<index_t> compute_counts_uniq(
    int64_t num_weights,
    const index_t* indices_data,
    int64_t indices_length,
    const std::vector<index_t>& counts) {
  std::vector<index_t> counts_uniq;
  counts_uniq.reserve(num_weights);
  int64_t o = 0;
  for (int64_t i = 0; i < indices_length; i += counts[indices_data[i]]) {
    counts_uniq.push_back(counts[indices_data[i]]);
    if (o > 0) {
      counts_uniq[o] += counts_uniq[o - 1];
    }
    o++;
  }
  return counts_uniq;
}