coalesced_unary_ufunc_out Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseUnaryOps.cpp lines 105–131

template <typename Ufunc>
Tensor& coalesced_unary_ufunc_out(const Tensor &self, Tensor &result, const Ufunc &ufunc) {
  if (self.is_same(result)) {
    TORCH_CHECK(self.is_coalesced(), "expected coalesced tensor for inplace operation");
    auto values = self._values();
    ufunc(values, values);
    return result;
  }

  TORCH_CHECK(self.is_sparse() && result.is_sparse());
  const auto input = self.coalesce();
  sparse_resize_(result, input.sizes(), input.sparse_dim(), input.dense_dim());
  auto *input_impl = sparse::get_sparse_impl(input);
  auto *result_impl = sparse::get_sparse_impl(result);

  auto input_values = input_impl->values();
  auto result_values = result_impl->values();
  result_values.resize_(input_values.sizes());
  ufunc(input_values, result_values);

  auto input_indices = input_impl->indices();
  auto result_indices = result_impl->indices();
  result_indices.resize_(input_indices.sizes());
  result_indices.copy_(input_indices);
  result._coalesced_(true);
  return result;
}