spmm_reduce_normalize_values_kernel_impl Class — pytorch Architecture

Source Code

aten/src/ATen/native/cpu/SpmmReduceKernel.cpp lines 347–373

template <typename scalar_t, typename index_t>
void spmm_reduce_normalize_values_kernel_impl(
    const Tensor& normalized_values,
    const Tensor& values,
    const Tensor& crow_indices,
    const Tensor& row_indices) {

  int64_t nnz = values.numel();
  if (nnz == 0) {
    return;
  }

  auto normalized_values_data = normalized_values.accessor<scalar_t, 1>();
  auto values_data = values.accessor<scalar_t, 1>();
  auto crow_data = crow_indices.accessor<index_t, 1>();
  auto row_data = row_indices.accessor<index_t, 1>();

  at::parallel_for(0, nnz, 1, [&](int64_t begin, int64_t end) {
    for (const auto i : c10::irange(begin, end)) {
      index_t row = row_data[i];
      index_t row_start = crow_data[row], row_end = crow_data[row + 1];
      // Note that when the row index row is listed in row_indices,
      // then crow_indices[row+1] > crow_indices[row] holds
      normalized_values_data[i] = values_data[i] / (row_end - row_start);
    }
  });
}