create_acc_buffer Class — pytorch Architecture

Source Code

aten/src/ATen/SparseCsrTensorUtils.h lines 425–446

template <typename acc_t, typename scalar_t>
inline std::tuple<Tensor, Tensor> create_acc_buffer(
    TensorOptions option,
    ScalarType type,
    int64_t nnz = -1) {
  Tensor new_values, new_values_acc;
  constexpr bool need_acc = !std::is_same_v<scalar_t, acc_t>;
  bool is_integral = at::isIntegralType(type, /*includeBool=*/true);
  if constexpr (need_acc) {
    auto acc_dtype = CppTypeToScalarType<acc_t>::value;
    new_values_acc = at::empty({}, option.dtype(acc_dtype));
    new_values = is_integral ? new_values_acc : at::empty({}, option);
  } else {
    new_values = new_values_acc = at::empty({}, option);
  }
  if (nnz != -1) {
    return std::make_tuple(
        new_values.resize_(nnz), new_values_acc.resize_(nnz));
  } else {
    return std::make_tuple(new_values, new_values_acc);
  }
}