1 #include "caffe2/operators/elementwise_ops_utils.h" 4 namespace elementwise_ops_utils {
6 std::tuple<size_t, size_t, size_t>
7 ComputeLegacyBroadcastSizes(
const Tensor&
A,
const Tensor&
B,
int axis) {
11 "If you are doing broadcasting, input1 should have " 12 "a smaller or equal number of dimensions.");
14 axis = A.dim() - B.dim();
17 axis >= 0 && axis <= A.dim() - B.dim(),
18 "Broadcast axis should be in the range of" 19 "[0, A.ndim() - B.ndim()], but axis = ",
23 while (b_dim_start < B.dim() && B.size(b_dim_start) == 1) {
26 int b_dim_end = B.dim() - 1;
27 while (b_dim_end >= b_dim_start && B.size(b_dim_end) == 1) {
30 size_t pre = 1, n = 1, post = 1;
31 for (
int i = 0; i < axis + b_dim_start; ++i) {
34 for (
int i = b_dim_start; i <= b_dim_end; ++i) {
36 A.size(i + axis), B.size(i),
"Broadcast dimension mismatch.");
39 for (
int i = axis + b_dim_end + 1; i < A.dim(); ++i) {
42 return std::make_tuple(pre, n, post);
45 std::vector<int> ComputeBinaryBroadcastForwardDims(
46 const std::vector<int>& A_dims,
47 const std::vector<int>& B_dims) {
48 const int ndim = std::max(A_dims.size(), B_dims.size());
49 std::vector<int> C_dims(ndim);
50 int i = A_dims.size() - 1;
51 int j = B_dims.size() - 1;
53 for (; i >= 0 && j >= 0; --k) {
54 const int A_dim = A_dims[i];
55 const int B_dim = B_dims[j];
56 CAFFE_ENFORCE(A_dim == B_dim || A_dim == 1 || B_dim == 1);
57 if (A_dim == 0 || B_dim == 0) {
60 C_dims[k] = std::max(A_dims[i], B_dims[j]);
66 C_dims[k--] = A_dims[i];
69 C_dims[k--] = B_dims[j];
74 void ComputeBinaryBroadcastBackwardAxes(
75 const std::vector<int>& A_dims,
76 const std::vector<int>& B_dims,
77 std::vector<int>* A_axes,
78 std::vector<int>* B_axes) {
81 const int ndim = std::max(A_dims.size(), B_dims.size());
82 int i = A_dims.size() - 1;
83 int j = B_dims.size() - 1;
85 for (; i >= 0 && j >= 0; --k) {
86 CAFFE_ENFORCE(A_dims[i] == B_dims[j] || A_dims[i] == 1 || B_dims[j] == 1);
87 if (A_dims[i] != B_dims[j]) {
100 A_axes->push_back(k);
103 for (; k >= 0; --k) {
104 B_axes->push_back(k);
107 std::reverse(A_axes->begin(), A_axes->end());
108 std::reverse(B_axes->begin(), B_axes->end());
111 void ComputeBinaryBroadcastBackwardDims(
112 const std::vector<int>& A_dims,
113 const std::vector<int>& B_dims,
114 std::vector<int>* A_back_dims,
115 std::vector<int>* B_back_dims) {
116 const int ndim = std::max(A_dims.size(), B_dims.size());
117 A_back_dims->assign(ndim, 1);
118 B_back_dims->assign(ndim, 1);
119 std::copy(A_dims.crbegin(), A_dims.crend(), A_back_dims->rbegin());
120 std::copy(B_dims.crbegin(), B_dims.crend(), B_back_dims->rbegin());
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...