tensor.hpp

Go to the documentation of this file.

#include <iostream>
#include <vector>

#ifndef DBG_TENSOR_HPP
#define DBG_TENSOR_HPP

template <typename T>
struct PrettyTensor final
{
    PrettyTensor (std::vector<uint16_t> datalimit) : datalimit_(datalimit) {}

    void print (std::ostream& out, T* arr, std::vector<uint8_t> shape)
    {
        print_helper(out, shape, arr, shape.size() - 1);
    }

    std::vector<uint16_t> datalimit_;

private:
    void print_helper (std::ostream& out, const std::vector<uint8_t>& shape,
        T* arr, uint8_t rank)
    {
        out << "[";
        uint16_t n = shape[rank];
        // apply limit only if limit is available
        if (rank < datalimit_.size())
        {
            n = std::min(n, datalimit_[rank]);
        }
        if (rank == 0)
        {
            out << arr[0];
            for (uint16_t i = 1; i < n; ++i)
            {
                out << "," << arr[i];
            }
        }
        else
        {
            auto it = shape.begin();
            size_t before = std::accumulate(it, it + rank, (size_t) 1,
                std::multiplies<size_t>());
            print_helper(out, shape, arr, rank - 1);
            for (uint16_t i = 1; i < n; ++i)
            {
                out << ",";
                print_helper(out, shape, arr + i * before, rank - 1);
            }
        }
        if (n < shape[rank])
        {
            out << "..";
        }
        out << "]";
    }
};

#endif // DBG_TENSOR_HPP