mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
c705d9eb1e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17991 changes: -Breaks bc: Tensor::type() now returns DeprecatedTypeProperties& rather than Type&. -Added DeprecatedTypeProperties, it serves as a temporary replacement for Type as the return value of Tensor::type(). This contributes to making Type just for dispatch purposes so that we can make it dtype agnostic. -Tensor::dispatch_type() now returns Type& like Tensor::type() used to do. -Changed callsites of Tensor::type() appropriately. Reviewed By: ezyang Differential Revision: D14443117 fbshipit-source-id: 239ccb7a09626279a71d1a37f8f82e7f57bf7d9e
77 lines
2.3 KiB
C++
77 lines
2.3 KiB
C++
#pragma once
|
|
|
|
#include <ATen/core/functional.h>
|
|
|
|
#include <ATen/ATen.h>
|
|
#include <utility>
|
|
|
|
namespace torch { namespace utils {
|
|
|
|
inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
|
|
static auto flatten = [](const at::Tensor &t) { return t.contiguous().view({-1}); };
|
|
if (tensors.size() == 1)
|
|
return flatten(tensors[0]);
|
|
return at::cat(fmap(tensors, flatten));
|
|
}
|
|
|
|
inline std::vector<at::Tensor> unflatten_dense_tensors(const at::Tensor& flat, at::TensorList tensors) {
|
|
std::vector<at::Tensor> outputs;
|
|
outputs.reserve(tensors.size());
|
|
size_t offset = 0;
|
|
for (const auto & tensor : tensors) {
|
|
auto numel = tensor.numel();
|
|
outputs.push_back(flat.narrow(0, offset, numel).view(tensor.sizes()));
|
|
offset += numel;
|
|
}
|
|
return outputs;
|
|
}
|
|
|
|
|
|
struct TensorGroup {
|
|
std::vector<at::Tensor> tensors;
|
|
size_t size = 0;
|
|
|
|
at::Type& type() {
|
|
AT_ASSERT(!tensors.empty());
|
|
return tensors[0].dispatch_type();
|
|
}
|
|
};
|
|
|
|
// Helper function that takes a list of tensors and splits them into tensor
|
|
// groups by the size limit and outputs these tensor groups. If the input
|
|
// tensors are of different tensor types, they will be split into different
|
|
// groups as well.
|
|
//
|
|
// Two options of splitting provided to the user,
|
|
//
|
|
// Imagine the size_limit is 256 and the list of input tensors are:
|
|
// tensor_a(fp16 - 128 bytes),
|
|
// tensor_b(fp32 - 256 bytes),
|
|
// tensor_c(fp16 - 128 bytes),
|
|
//
|
|
// when fine_grained == false:
|
|
// The function will read the list of tensors sequentially and accumulate
|
|
// enough tensors for each data type until the size_limit, therefore:
|
|
// it will output: {{tensor_a, tensor_c}, {tensor_b}}
|
|
//
|
|
// when fine_grained == true:
|
|
// The function will read the list of tensors sequentially and accumulate
|
|
// enough tensors for all data types until the size_limit, and then split
|
|
// the accumulated tensors into different groups by data types, therefore:
|
|
// it will output: {{tensor_a}, {tensor_b}, {tensor_c}}
|
|
std::vector<TensorGroup> take_tensors(
|
|
at::TensorList tensors,
|
|
size_t size_limit,
|
|
bool fine_grained = false);
|
|
|
|
void reorder_tensors_like(std::vector<at::Tensor>& tensors, at::TensorList order);
|
|
|
|
std::pair<at::Tensor, at::Tensor> flatten_sparse_tensors(at::TensorList tensors);
|
|
|
|
std::vector<at::Tensor> unflatten_sparse_tensors(
|
|
const at::Tensor& flat_indices,
|
|
const at::Tensor& flat_values,
|
|
at::TensorList tensors);
|
|
|
|
}}
|