| 252 | } |
| 253 | |
| 254 | Maybe<Shape> InferShapeUnspecifiedDim(const int64_t& elem_count, const Shape& shape) { |
| 255 | int need_infer_axis = -1; |
| 256 | int64_t target_elem_count = 1; |
| 257 | for (int i = 0; i < shape.NumAxes(); ++i) { |
| 258 | if (shape.At(i) < -1) { |
| 259 | return Error::RuntimeError() << "Invalid shape dimension " << shape.At(i); |
| 260 | } else if (shape.At(i) == -1) { |
| 261 | CHECK_OR_RETURN_ERROR(need_infer_axis == -1) |
| 262 | << Error::RuntimeError() << "only one dimension can be inferred"; |
| 263 | need_infer_axis = i; |
| 264 | } else { |
| 265 | target_elem_count *= shape.At(i); |
| 266 | } |
| 267 | } |
| 268 | Shape infered_shape = shape; |
| 269 | if (need_infer_axis == -1) { |
| 270 | if (elem_count > 0) { |
| 271 | // For 0-size tensor, we don't need to check the element size. |
| 272 | CHECK_OR_RETURN_ERROR(target_elem_count == elem_count) |
| 273 | << Error::RuntimeError() << "shape '" << shape.ToString() |
| 274 | << "' is invalid for input of size " << elem_count; |
| 275 | } |
| 276 | } else { |
| 277 | infered_shape.Set(need_infer_axis, elem_count / target_elem_count); |
| 278 | CHECK_OR_RETURN_ERROR(target_elem_count * infered_shape.At(need_infer_axis) == elem_count) |
| 279 | << Error::RuntimeError() << "shape '" << shape.ToString() |
| 280 | << "' is invalid for input of size " << elem_count; |
| 281 | } |
| 282 | return infered_shape; |
| 283 | } |
| 284 | |
| 285 | Maybe<Shape> InferUnifiedShapeForBroadcasting(const std::vector<Shape>& shapes) { |
| 286 | if (shapes.empty()) { return Error::RuntimeError() << "shapes should not be empty."; } |
no test coverage detected