1 // This file is part of Eigen, a lightweight C++ template library 2 // for linear algebra. 3 // 4 // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog (at) gmail.com> 5 // 6 // This Source Code Form is subject to the terms of the Mozilla 7 // Public License v. 2.0. If a copy of the MPL was not distributed 8 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/. 9 10 #include "main.h" 11 12 #include <Eigen/CXX11/Tensor> 13 14 using Eigen::Tensor; 15 16 17 struct InsertZeros { 18 DSizes<DenseIndex, 2> dimensions(const Tensor<float, 2>& input) const { 19 DSizes<DenseIndex, 2> result; 20 result[0] = input.dimension(0) * 2; 21 result[1] = input.dimension(1) * 2; 22 return result; 23 } 24 25 template <typename Output, typename Device> 26 void eval(const Tensor<float, 2>& input, Output& output, const Device& device) const 27 { 28 array<DenseIndex, 2> strides; 29 strides[0] = 2; 30 strides[1] = 2; 31 output.stride(strides).device(device) = input; 32 33 Eigen::DSizes<DenseIndex, 2> offsets(1,1); 34 Eigen::DSizes<DenseIndex, 2> extents(output.dimension(0)-1, output.dimension(1)-1); 35 output.slice(offsets, extents).stride(strides).device(device) = input.constant(0.0f); 36 } 37 }; 38 39 static void test_custom_unary_op() 40 { 41 Tensor<float, 2> tensor(3,5); 42 tensor.setRandom(); 43 44 Tensor<float, 2> result = tensor.customOp(InsertZeros()); 45 VERIFY_IS_EQUAL(result.dimension(0), 6); 46 VERIFY_IS_EQUAL(result.dimension(1), 10); 47 48 for (int i = 0; i < 6; i+=2) { 49 for (int j = 0; j < 10; j+=2) { 50 VERIFY_IS_EQUAL(result(i, j), tensor(i/2, j/2)); 51 } 52 } 53 for (int i = 1; i < 6; i+=2) { 54 for (int j = 1; j < 10; j+=2) { 55 VERIFY_IS_EQUAL(result(i, j), 0); 56 } 57 } 58 } 59 60 61 struct BatchMatMul { 62 DSizes<DenseIndex, 3> dimensions(const Tensor<float, 3>& input1, const Tensor<float, 3>& input2) const { 63 DSizes<DenseIndex, 3> result; 64 result[0] = input1.dimension(0); 65 result[1] = input2.dimension(1); 66 result[2] = input2.dimension(2); 67 return result; 68 } 69 70 template <typename Output, typename Device> 71 void eval(const Tensor<float, 3>& input1, const Tensor<float, 3>& input2, 72 Output& output, const Device& device) const 73 { 74 typedef Tensor<float, 3>::DimensionPair DimPair; 75 array<DimPair, 1> dims; 76 dims[0] = DimPair(1, 0); 77 for (int i = 0; i < output.dimension(2); ++i) { 78 output.template chip<2>(i).device(device) = input1.chip<2>(i).contract(input2.chip<2>(i), dims); 79 } 80 } 81 }; 82 83 84 static void test_custom_binary_op() 85 { 86 Tensor<float, 3> tensor1(2,3,5); 87 tensor1.setRandom(); 88 Tensor<float, 3> tensor2(3,7,5); 89 tensor2.setRandom(); 90 91 Tensor<float, 3> result = tensor1.customOp(tensor2, BatchMatMul()); 92 for (int i = 0; i < 5; ++i) { 93 typedef Tensor<float, 3>::DimensionPair DimPair; 94 array<DimPair, 1> dims; 95 dims[0] = DimPair(1, 0); 96 Tensor<float, 2> reference = tensor1.chip<2>(i).contract(tensor2.chip<2>(i), dims); 97 TensorRef<Tensor<float, 2> > val = result.chip<2>(i); 98 for (int j = 0; j < 2; ++j) { 99 for (int k = 0; k < 7; ++k) { 100 VERIFY_IS_APPROX(val(j, k), reference(j, k)); 101 } 102 } 103 } 104 } 105 106 107 void test_cxx11_tensor_custom_op() 108 { 109 CALL_SUBTEST(test_custom_unary_op()); 110 CALL_SUBTEST(test_custom_binary_op()); 111 } 112
