| 22 | void BCast::Reverse(Vec* shape) { std::reverse(shape->begin(), shape->end()); } |
| 23 | |
| 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; |