FM / utils

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/utils/

/usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/utils

Name	Size	Mode	Actions
auto_gil.h	1034	0644	edit dl rm
byte_order.h	2484	0644	edit dl rm
crash_handler.h	1146	0644	edit dl rm
cuda_enabled.h	154	0644	edit dl rm
cuda_lazy_init.h	974	0644	edit dl rm
disable_torch_function.h	863	0644	edit dl rm
disallow_copy.h	103	0644	edit dl rm
init.h	324	0644	edit dl rm
invalid_arguments.h	302	0644	edit dl rm
memory.h	1173	0644	edit dl rm
numpy_stub.h	399	0644	edit dl rm
object_ptr.h	1329	0644	edit dl rm
out_types.h	294	0644	edit dl rm
pybind.h	7966	0644	edit dl rm
pycfunction_helpers.h	209	0644	edit dl rm
python_arg_parser.h	30753	0644	edit dl rm
python_compat.h	2904	0644	edit dl rm
python_dispatch.h	174	0644	edit dl rm
python_numbers.h	5064	0644	edit dl rm
python_scalars.h	2928	0644	edit dl rm
python_strings.h	4592	0644	edit dl rm
python_stub.h	56	0644	edit dl rm
python_tuples.h	684	0644	edit dl rm
six.h	1425	0644	edit dl rm
structseq.h	153	0644	edit dl rm
tensor_apply.h	431	0644	edit dl rm
tensor_dtypes.h	244	0644	edit dl rm
tensor_flatten.h	2780	0644	edit dl rm
tensor_layouts.h	107	0644	edit dl rm
tensor_list.h	196	0644	edit dl rm
tensor_memoryformats.h	113	0644	edit dl rm
tensor_new.h	1807	0644	edit dl rm
tensor_numpy.h	542	0644	edit dl rm
tensor_qschemes.h	186	0644	edit dl rm
tensor_types.h	491	0644	edit dl rm
throughput_benchmark-inl.h	5210	0644	edit dl rm
throughput_benchmark.h	6868	0644	edit dl rm
variadic.h	4394	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/torch/csrc/utils/tensor_flatten.h (2780B)

#pragma once

#include <ATen/core/functional.h>
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <ATen/ATen.h>
#include <utility>
#include <c10/core/TensorOptions.h>

namespace torch { namespace utils {

/// Generate an ID for a combination of tensor backend + scalar type to be used
/// when ordering tensors ('like' tensors are grouped by pulling out their
/// backend + scalar type, so this function combines that into a single number)
inline size_t type_id(const at::Tensor& tensor) {
  return static_cast<size_t>(tensor.options().backend()) *
      static_cast<size_t>(at::ScalarType::NumOptions) +
      static_cast<size_t>(tensor.scalar_type());
}

inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
  return at::flatten_dense_tensors(tensors);
}

inline std::vector<at::Tensor> unflatten_dense_tensors(const at::Tensor& flat, at::TensorList tensors) {
  return at::unflatten_dense_tensors(flat, tensors);
}

// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
struct TensorGroup {
  std::vector<at::Tensor> tensors;
  size_t size = 0;

size_t type_id() {
    AT_ASSERT(!tensors.empty());
    return ::torch::utils::type_id(tensors[0]);
  }

const at::TensorOptions options() {
    AT_ASSERT(!tensors.empty());
    return tensors[0].options();
  }
};

// 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}}
TORCH_API std::vector<TensorGroup> take_tensors(
    at::TensorList tensors,
    size_t size_limit,
    bool fine_grained = false);

TORCH_API void reorder_tensors_like(std::vector<at::Tensor>& tensors, at::TensorList order);

TORCH_API std::pair<at::Tensor, at::Tensor> flatten_sparse_tensors(at::TensorList tensors);

TORCH_API std::vector<at::Tensor> unflatten_sparse_tensors(
    const at::Tensor& flat_indices,
    const at::Tensor& flat_values,
    at::TensorList tensors);

}}

cmd: