1 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/util/bcast.h"
17
18 #include "tensorflow/core/platform/logging.h"
19 namespace tensorflow {
20
21 /* static */
Reverse(Vec * shape)22 void BCast::Reverse(Vec* shape) { std::reverse(shape->begin(), shape->end()); }
23
BCast(const Vec & sx,const Vec & sy,const bool fewer_dims_optimization)24 BCast::BCast(const Vec& sx, const Vec& sy, const bool fewer_dims_optimization) {
25 if (sx == sy && TF_PREDICT_TRUE(fewer_dims_optimization)) {
26 // Fast path for common case of identical shapes for sx and sy
27 int64 elements = 1;
28 const int n = sx.size();
29 output_.resize(n);
30 for (int i = 0; i < n; i++) {
31 const int64 dim = sx[i];
32 elements *= dim;
33 output_[i] = dim;
34 }
35 result_.push_back(elements);
36 x_reshape_.push_back(elements);
37 y_reshape_.push_back(elements);
38 x_bcast_.push_back(1);
39 y_bcast_.push_back(1);
40 // grad_x_reduce_ and grad_y_reduce_ are left as empty
41 } else {
42 // Reverse the shape of x and y for convenience.
43 // After the reverse, 0-th is the inner-most dimension.
44 Vec x = sx;
45 Vec y = sy;
46 Reverse(&x);
47 Reverse(&y);
48
49 // 1-extend and align x and y so that they are the same size.
50 if (x.size() > y.size()) {
51 y.resize(x.size(), 1);
52 } else {
53 x.resize(y.size(), 1);
54 }
55
56 // Going through each dimension starting from the inner-most
57 // dimension, compares dimension of x and y. They are compatible if
58 // they are equal or either is 1.
59 enum State {
60 UNKNOWN,
61 SAME,
62 X_ONE,
63 Y_ONE,
64 };
65 State prev = UNKNOWN;
66 const int64 n = x.size();
67 for (int i = 0; i < n; ++i) {
68 // Output shape.
69 State curr = UNKNOWN;
70 const int64 x_i = x[i]; // i-th dimension of x.
71 const int64 y_i = y[i]; // i-th dimension of y.
72 int64 o_i; // i-th dimension of the output.
73 int64 bx_i; // i-th broadcast for x.
74 int64 by_i; // i-th broadcast for y.
75 // Invariant:
76 // o_i = x_i * bx_i = y_i * by_i
77 if (x_i == y_i) {
78 // No broadcast.
79 o_i = x_i;
80 bx_i = 1;
81 by_i = 1;
82 curr = SAME;
83 } else if (x_i == 1) {
84 // x broadcast to y on this dimension.
85 o_i = y_i;
86 bx_i = y_i;
87 by_i = 1;
88 grad_x_reduce_idx_.push_back(n - 1 - i);
89 curr = X_ONE;
90 } else if (y_i == 1) {
91 // y broadcast to x on this dimension.
92 o_i = x_i;
93 bx_i = 1;
94 by_i = x_i;
95 grad_y_reduce_idx_.push_back(n - 1 - i);
96 curr = Y_ONE;
97 } else {
98 valid_ = false;
99 return;
100 }
101 output_.push_back(o_i);
102 // Reshape/broadcast.
103 // Invariant:
104 // result[i] == x_reshape[i] * x_bcast[i] == y_reshape_[i] * y_bcast_[i]
105 if (curr == SAME && x_i == 1) {
106 // Both side are 1s.
107 grad_x_reduce_idx_.push_back(n - 1 - i);
108 grad_y_reduce_idx_.push_back(n - 1 - i);
109 if (!TF_PREDICT_TRUE(fewer_dims_optimization)) {
110 result_.push_back(o_i);
111 x_reshape_.push_back(x_i);
112 x_bcast_.push_back(bx_i);
113 y_reshape_.push_back(y_i);
114 y_bcast_.push_back(by_i);
115 }
116 continue;
117 } else if (TF_PREDICT_TRUE(fewer_dims_optimization) && prev == curr) {
118 // It is a run of the same cases(no broadcast, x broadcast to y, y
119 // broadcast to x). We can reshape the input so that fewer dimensions
120 // are involved in the intermediate computation.
121 result_.back() *= o_i;
122 x_reshape_.back() *= x_i;
123 x_bcast_.back() *= bx_i;
124 y_reshape_.back() *= y_i;
125 y_bcast_.back() *= by_i;
126 } else {
127 result_.push_back(o_i);
128 x_reshape_.push_back(x_i);
129 x_bcast_.push_back(bx_i);
130 y_reshape_.push_back(y_i);
131 y_bcast_.push_back(by_i);
132 }
133 prev = curr;
134 }
135
136 if (result_.empty()) {
137 // Can happen when both x and y are effectively scalar.
138 result_.push_back(1);
139 x_reshape_.push_back(1);
140 x_bcast_.push_back(1);
141 y_reshape_.push_back(1);
142 y_bcast_.push_back(1);
143 }
144
145 // Reverse all vectors since x and y were reversed at very
146 // beginning.
147 Reverse(&x_reshape_);
148 Reverse(&x_bcast_);
149 Reverse(&y_reshape_);
150 Reverse(&y_bcast_);
151 Reverse(&result_);
152 Reverse(&output_);
153 Reverse(&grad_x_reduce_idx_);
154 Reverse(&grad_y_reduce_idx_);
155 }
156 }
157
FromShape(const TensorShape & shape)158 BCast::Vec BCast::FromShape(const TensorShape& shape) {
159 const int N = shape.dims();
160 BCast::Vec ret(N);
161 for (int i = 0; i < N; ++i) {
162 ret[i] = shape.dim_size(i);
163 }
164 return ret;
165 }
166
ToShape(const BCast::Vec & vec)167 TensorShape BCast::ToShape(const BCast::Vec& vec) {
168 TensorShape shape(vec);
169 return shape;
170 }
171
172 } // end namespace tensorflow
173