static */
| 354 | } |
| 355 | |
| 356 | /* static */ StatusOr<Shape> ShapeInference::InferConcatOpShape( |
| 357 | absl::Span<const Shape* const> arg_shapes, const int64 dimension) { |
| 358 | if (arg_shapes.empty()) { |
| 359 | return InvalidArgument("Concatenate expects at least one argument."); |
| 360 | } |
| 361 | if (dimension < 0 || dimension >= arg_shapes[0]->rank()) { |
| 362 | return InvalidArgument("Concatenate dimension out of bounds: %d.", |
| 363 | dimension); |
| 364 | } |
| 365 | const Shape* arg_shape = nullptr; |
| 366 | PrimitiveType element_type = PRIMITIVE_TYPE_INVALID; |
| 367 | for (const Shape* shape : arg_shapes) { |
| 368 | TF_RETURN_IF_ERROR(ExpectArray(*shape, "operand of concatenation")); |
| 369 | if (!arg_shape) { |
| 370 | arg_shape = shape; |
| 371 | element_type = arg_shape->element_type(); |
| 372 | continue; |
| 373 | } |
| 374 | if (arg_shape->rank() != shape->rank()) { |
| 375 | return InvalidArgument( |
| 376 | "Cannot concatenate arrays with different ranks: %d (%s) vs %d " |
| 377 | "(%s).", |
| 378 | arg_shape->rank(), ShapeUtil::HumanString(*arg_shape), shape->rank(), |
| 379 | ShapeUtil::HumanString(*shape)); |
| 380 | } |
| 381 | if (!ShapeUtil::SameElementTypeIgnoringFpPrecision(*arg_shape, *shape)) { |
| 382 | return InvalidArgument( |
| 383 | "Cannot concatenate arrays with different element types: %s vs %s.", |
| 384 | PrimitiveType_Name(arg_shape->element_type()), |
| 385 | PrimitiveType_Name(shape->element_type())); |
| 386 | } |
| 387 | for (int64 dimension_number = 0; dimension_number < arg_shape->rank(); |
| 388 | ++dimension_number) { |
| 389 | if (arg_shape->dimensions(dimension_number) != |
| 390 | shape->dimensions(dimension_number)) { |
| 391 | if (dimension_number == dimension) { |
| 392 | continue; // It's okay to differ in the dimension we're |
| 393 | // concatenating. |
| 394 | } |
| 395 | return InvalidArgument( |
| 396 | "Cannot concatenate arrays that differ in dimensions other than " |
| 397 | "the one being concatenated (the other array dimensions must be " |
| 398 | "the same): %s vs %s in dimension %d.", |
| 399 | ShapeUtil::HumanString(*arg_shape), ShapeUtil::HumanString(*shape), |
| 400 | dimension); |
| 401 | } |
| 402 | } |
| 403 | element_type = ShapeUtil::HigherPrecisionElementType(*shape, *arg_shape); |
| 404 | } |
| 405 | |
| 406 | std::vector<int64> new_dimensions(arg_shape->dimensions().begin(), |
| 407 | arg_shape->dimensions().end()); |
| 408 | for (size_t i = 1; i < arg_shapes.size(); ++i) { |
| 409 | new_dimensions[dimension] += arg_shapes[i]->dimensions(dimension); |
| 410 | } |
| 411 | |
| 412 | Shape result = ShapeUtil::MakeShape(element_type, new_dimensions); |
| 413 |
nothing calls this directly
no test coverage detected