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@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 {
dimensionsInsertZeros18 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>
evalInsertZeros26 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
test_custom_unary_op()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 {
dimensionsBatchMatMul62 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>
evalBatchMatMul71 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
test_custom_binary_op()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
test_cxx11_tensor_custom_op()107 void test_cxx11_tensor_custom_op()
108 {
109 CALL_SUBTEST(test_custom_unary_op());
110 CALL_SUBTEST(test_custom_binary_op());
111 }
112