Added ability to print a tensor using an iostream.

This commit is contained in:
Benoit Steiner 2014-10-10 16:17:26 -07:00
parent 2ed1838aeb
commit 0219f8aed4
4 changed files with 117 additions and 0 deletions

View File

@ -64,6 +64,8 @@
#include "unsupported/Eigen/CXX11/src/Tensor/TensorFixedSize.h"
#include "unsupported/Eigen/CXX11/src/Tensor/TensorMap.h"
#include "unsupported/Eigen/CXX11/src/Tensor/TensorIO.h"
#include "Eigen/src/Core/util/ReenableStupidWarnings.h"
#endif // EIGEN_CXX11_TENSOR_MODULE

View File

@ -0,0 +1,44 @@
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#ifndef EIGEN_CXX11_TENSOR_TENSOR_IO_H
#define EIGEN_CXX11_TENSOR_TENSOR_IO_H
namespace Eigen {
template <typename T>
std::ostream& operator << (std::ostream& os, const TensorBase<T, ReadOnlyAccessors>& expr) {
// Evaluate the expression if needed
TensorForcedEvalOp<const T> eval = expr.eval();
TensorEvaluator<const TensorForcedEvalOp<const T>, DefaultDevice> tensor(eval, DefaultDevice());
tensor.evalSubExprsIfNeeded(NULL);
typedef typename T::Scalar Scalar;
typedef typename T::Index Index;
typedef typename TensorEvaluator<const TensorForcedEvalOp<const T>, DefaultDevice>::Dimensions Dimensions;
const Index total_size = internal::array_prod(tensor.dimensions());
// Print the tensor as a 1d vector or a 2d matrix.
if (internal::array_size<Dimensions>::value == 1) {
Map<Array<Scalar, Dynamic, 1> > array(tensor.data(), total_size);
os << array;
} else {
const Index first_dim = tensor.dimensions()[0];
Map<Array<Scalar, Dynamic, Dynamic> > matrix(tensor.data(), first_dim, total_size/first_dim);
os << matrix;
}
// Cleanup.
tensor.cleanup();
return os;
}
} // end namespace Eigen
#endif // EIGEN_CXX11_TENSOR_TENSOR_IO_H

View File

@ -124,4 +124,5 @@ if(EIGEN_TEST_CXX11)
ei_add_test(cxx11_tensor_striding "-std=c++0x")
# ei_add_test(cxx11_tensor_device "-std=c++0x")
ei_add_test(cxx11_tensor_thread_pool "-std=c++0x")
ei_add_test(cxx11_tensor_io "-std=c++0x")
endif()

View File

@ -0,0 +1,70 @@
// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
#include "main.h"
#include <sstream>
#include <string>
#include <Eigen/CXX11/Tensor>
static void test_output_1d()
{
Tensor<int, 1> tensor(5);
for (int i = 0; i < 5; ++i) {
tensor(i) = i;
}
std::stringstream os;
os << tensor;
std::string expected("0\n1\n2\n3\n4");
VERIFY_IS_EQUAL(std::string(os.str()), expected);
}
static void test_output_2d()
{
Tensor<int, 2> tensor(5, 3);
for (int i = 0; i < 5; ++i) {
for (int j = 0; j < 3; ++j) {
tensor(i, j) = i*j;
}
}
std::stringstream os;
os << tensor;
std::string expected("0 0 0\n0 1 2\n0 2 4\n0 3 6\n0 4 8");
VERIFY_IS_EQUAL(std::string(os.str()), expected);
}
static void test_output_expr()
{
Tensor<int, 1> tensor1(5);
Tensor<int, 1> tensor2(5);
for (int i = 0; i < 5; ++i) {
tensor1(i) = i;
tensor2(i) = 7;
}
std::stringstream os;
os << tensor1 + tensor2;
std::string expected(" 7\n 8\n 9\n10\n11");
VERIFY_IS_EQUAL(std::string(os.str()), expected);
}
void test_cxx11_tensor_io()
{
CALL_SUBTEST(test_output_1d());
CALL_SUBTEST(test_output_2d());
CALL_SUBTEST(test_output_expr());
}