| 171 | } |
| 172 | |
| 173 | std::tuple<bool, Shape, std::vector<Strides>> compiled_collapse_contiguous_dims( |
| 174 | const std::vector<array>& inputs, |
| 175 | const array& out, |
| 176 | const std::function<bool(size_t)>& is_constant) { |
| 177 | const Shape& shape = out.shape(); |
| 178 | bool contiguous = compiled_check_contiguity(inputs, shape); |
| 179 | if (contiguous) { |
| 180 | return {true, shape, {}}; |
| 181 | } |
| 182 | |
| 183 | std::vector<Strides> strides_vec{out.strides()}; |
| 184 | for (size_t i = 0; i < inputs.size(); ++i) { |
| 185 | // Skip constants. |
| 186 | if (is_constant(i)) { |
| 187 | continue; |
| 188 | } |
| 189 | |
| 190 | // Skip scalar inputs. |
| 191 | const auto& x = inputs[i]; |
| 192 | if (is_scalar(x)) { |
| 193 | continue; |
| 194 | } |
| 195 | |
| 196 | // Broadcast the inputs to the output shape. |
| 197 | Strides xstrides; |
| 198 | size_t j = 0; |
| 199 | for (; j < shape.size() - x.ndim(); ++j) { |
| 200 | if (shape[j] == 1) { |
| 201 | xstrides.push_back(out.strides()[j]); |
| 202 | } else { |
| 203 | xstrides.push_back(0); |
| 204 | } |
| 205 | } |
| 206 | for (size_t i = 0; i < x.ndim(); ++i, ++j) { |
| 207 | if (x.shape(i) == 1) { |
| 208 | if (shape[j] == 1) { |
| 209 | xstrides.push_back(out.strides()[j]); |
| 210 | } else { |
| 211 | xstrides.push_back(0); |
| 212 | } |
| 213 | } else { |
| 214 | xstrides.push_back(x.strides()[i]); |
| 215 | } |
| 216 | } |
| 217 | strides_vec.push_back(std::move(xstrides)); |
| 218 | } |
| 219 | |
| 220 | auto tup = collapse_contiguous_dims(shape, strides_vec, INT32_MAX); |
| 221 | return {false, std::move(std::get<0>(tup)), std::move(std::get<1>(tup))}; |
| 222 | } |
| 223 | |
| 224 | bool compiled_use_large_index( |
| 225 | const std::vector<array>& inputs, |
no test coverage detected