validate_compressed_sparse_indices_kernel Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/ValidateCompressedIndicesCommon.h lines 361–410

template <
    template <typename func_t>
    class kernel_t,
    template <typename func_t, typename vec_func_t>
    class vec_kernel_t = EmptyVecKernel,
    template <typename scalar_t> class Vec = DummyVec>
void validate_compressed_sparse_indices_kernel(
    const bool is_crow,
    const Tensor& cidx,
    const Tensor& idx,
    const int64_t cdim,
    const int64_t dim,
    const int64_t nnz) {
  constexpr size_t idx_max_ndims = 8; // up to 7-dim batch.
  const size_t idx_ndims = static_cast<size_t>(idx.dim());

  if (is_crow) {
    if (idx_ndims <= idx_max_ndims) {
      _validate_compressed_sparse_indices_kernel<
          CDimName::CRow,
          kernel_t,
          vec_kernel_t,
          Vec,
          idx_max_ndims>(cidx, idx, cdim, dim, nnz);
    }
    else {
      _validate_compressed_sparse_indices_kernel<
          CDimName::CRow,
          kernel_t,
          vec_kernel_t,
          Vec>(cidx, idx, cdim, dim, nnz);
    }
  } else {
    if (idx_ndims <= idx_max_ndims) {
      _validate_compressed_sparse_indices_kernel<
          CDimName::CCol,
          kernel_t,
          vec_kernel_t,
          Vec,
          idx_max_ndims>(cidx, idx, cdim, dim, nnz);
    }
    else {
      _validate_compressed_sparse_indices_kernel<
          CDimName::CCol,
          kernel_t,
          vec_kernel_t,
          Vec>(cidx, idx, cdim, dim, nnz);
    }
  }
}