required_layout Class — pytorch Architecture

Source Code

aten/src/ATen/native/sparse/SparseCsrTensor.cpp lines 485–507

template <Layout required_layout>
static Tensor _sparse_compressed_tensor_unsafe_template(const Tensor& compressed_indices,
                                                 const Tensor& plain_indices,
                                                 const Tensor& values,
                                                 IntArrayRef size,
                                                 std::optional<ScalarType> dtype,
                                                 std::optional<Layout> layout,
                                                 std::optional<Device> device,
                                                 std::optional<bool> pin_memory) {
  Layout layout_ = layout.value_or(required_layout);
  TORCH_CHECK(layout_ == required_layout, "sparse compressed layout must be ",required_layout, " but got ", layout_);
  if (at::globalContext().checkSparseTensorInvariants()) {
    _validate_sparse_compressed_tensor_args_worker(compressed_indices, plain_indices, values, size, layout_, true);
  }
  TensorOptions options = TensorOptions().dtype(dtype).layout(layout_).device(device).pinned_memory(pin_memory);
  SparseCsrTensor self = new_compressed_tensor(options);
  if (pin_memory.value_or(false) && !values.is_pinned()) {
    get_sparse_csr_impl(self)->set_member_tensors(compressed_indices.pin_memory(), plain_indices.pin_memory(), values.pin_memory(), size);
  } else {
    get_sparse_csr_impl(self)->set_member_tensors(compressed_indices, plain_indices, values, size);
  }
  return self;
}