tensor_utils.h

Go to the documentation of this file.

#ifndef ATEN_CORE_TENSOR_UTILS_H_
#define ATEN_CORE_TENSOR_UTILS_H_
 
#include <ATen/core/tensor.h>
#include <ATen/core/tensor_shape.h>
 
namespace container {
 
__inline__
std::string removeTrailingZeros(std::string str) {
    int i = static_cast<int>(str.length()) - 1;
    while (i >= 0 && str[i] == '0') {
        i--;
    }
    if (i == -1) {
        return "0";
    }
    return str.substr(0, i + 1);
}
 
template<typename T>
__inline__
int _get_digit_places(
        const T* arr,
        int size,
        int& integer_count,
        int& fraction_count)
{
    integer_count = 0;
    fraction_count = 0;
 
    for (int i = 0; i < size; i++) {
        int digits = 0;
        if (arr[i] < 0) {
            digits = log10(-arr[i]) + 1;
            if (digits + 1 > integer_count) {
                integer_count = digits + 1;
            }
        }
        else {
            digits = log10(arr[i]) + 1;
            if (digits > integer_count) {
                integer_count = digits;
            }
        }
 
        T fraction = arr[i] - std::floor(arr[i]);
        if (fraction == 0) {
            continue;
        }
        std::string str = removeTrailingZeros(std::to_string(fraction));
        digits = str.length() - str.find('.');
        if (digits > fraction_count) {
            fraction_count = digits;
        }
    }
 
    return integer_count + fraction_count;
}
 
template<typename T>
__inline__
int _get_digit_places(
        const std::complex<T>* arr,
        int size,
        int& integer_count,
        int& fraction_count)
{
    return _get_digit_places<T>(reinterpret_cast<const T*>(arr), size * 2, integer_count, fraction_count);
}
 
template <typename T>
__inline__
void _output_wrapper(
        std::ostream& os,
        const T data,
        const int& digit_width,
        const int& fraction_count)
{
    os << std::setw(digit_width) \
       << std::setprecision(fraction_count) << std::fixed << data;
}
 
template <typename T>
__inline__
void _output_wrapper(
        std::ostream& os,
        const std::complex<T> data,
        const int& digit_width,
        const int& fraction_count)
{
    // Write the real and imaginary parts of the complex value to the output stream
    // with the specified formatting.
    os << "{";
    os << std::setw(digit_width) \
       << std::setprecision(fraction_count) << std::fixed
       << data.real();
    os << ", ";
    os << std::setw(digit_width) \
       << std::setprecision(fraction_count) << std::fixed
       << data.imag();
    os << "}";
}
 
template <>
__inline__
void _output_wrapper(
        std::ostream& os,
        const int data,
        const int& digit_width,
        const int& fraction_count)
{
    os << std::setw(digit_width - 1) \
       << std::setprecision(fraction_count) << std::fixed << data;
}
 
template <typename T>
__inline__
void _internal_output(
        std::ostream& os,
        const T * data,
        const TensorShape& shape,
        const int64_t& num_elements)
{
    int integer_count = 0, fraction_count = 0;
    int digit_width = _get_digit_places(data, num_elements, integer_count, fraction_count) + 1;
    if (shape.ndim() == 1) {
        os << "[";
        for (int i = 0; i < num_elements; ++i) {
            _output_wrapper(os, data[i], digit_width, fraction_count);
            if (i != num_elements - 1) {
                os << ",";
            }
        }
        os << "]";
    }
    else if (shape.ndim() == 2) {
        os << "[";
        for (int i = 0; i < shape.dim_size(0); ++i) {
            if (i != 0) os << "       ";
            os << "[";
            for (int j = 0; j < shape.dim_size(1); ++j) {
                _output_wrapper(os, data[i * shape.dim_size(1) + j], digit_width, fraction_count);
                if (j != shape.dim_size(1) - 1) {
                    os << ", ";
                }
            }
            os << "]";
            if (i != shape.dim_size(0) - 1) os << ",\n";
        }
        os << "]";
    }
    else if (shape.ndim() == 3) {
        os << "[";
        for (int i = 0; i < shape.dim_size(0); ++i) {
            if (i != 0) os << "       ";
            os << "[";
            for (int j = 0; j < shape.dim_size(1); ++j) {
                if (j != 0) os << "        ";
                os << "[";
                for (int k = 0; k < shape.dim_size(2); ++k) {
                    _output_wrapper(os, data[i * shape.dim_size(1) * shape.dim_size(2) + j * shape.dim_size(2) + k], digit_width, fraction_count);
                    if (k != shape.dim_size(2) - 1) {
                        os << ", ";
                    }
                }
                os << "]";
                if (j != shape.dim_size(1) - 1) os << ",\n";
            }
            os << "]";
            if (i != shape.dim_size(0) - 1) os << ",\n\n";
        }
        os << "]";
    }
    else {
        for (int64_t i = 0; i < num_elements; ++i) {
            _output_wrapper(os, data[i], 0, 0);
            if (i < num_elements - 1) {
                os << ", ";
            }
        }
    }
}
 
template <typename T>
T extract(const container::Tensor& tensor) {
    if (tensor.device_type() == DeviceType::CpuDevice) {
        return reinterpret_cast<T*>(tensor.data())[0];
    }
    else {
        T result = 0;
        TEMPLATE_ALL_2(tensor.data_type(), tensor.device_type(),
            kernels::synchronize_memory<T, DEVICE_CPU, DEVICE_>()(
                &result, reinterpret_cast<T*>(tensor.data()), 1))
        return result;
    }
}
 
} // namespace container
 
#endif // ATEN_CORE_TENSOR_UTILS_H_