| 486 | } |
| 487 | |
| 488 | Status InferenceContext::MergePrefix(ShapeHandle s, ShapeHandle prefix, |
| 489 | ShapeHandle* s_out, |
| 490 | ShapeHandle* prefix_out) { |
| 491 | *s_out = *prefix_out = nullptr; |
| 492 | if (!RankKnown(prefix) || !RankKnown(s)) { |
| 493 | *s_out = s; |
| 494 | *prefix_out = prefix; |
| 495 | return Status::OK(); |
| 496 | } |
| 497 | const int32 rank = Rank(prefix); |
| 498 | TF_RETURN_IF_ERROR(WithRankAtLeast(s, rank, &s)); |
| 499 | |
| 500 | // Merge the prefix dims and create the new output shapes. |
| 501 | const int32 rank_s = Rank(s); |
| 502 | std::vector<DimensionHandle> dims; |
| 503 | dims.reserve(std::max(rank, rank_s)); |
| 504 | dims.resize(rank); |
| 505 | for (int i = 0; i < rank; ++i) { |
| 506 | TF_RETURN_IF_ERROR(Merge(Dim(s, i), Dim(prefix, i), &dims[i])); |
| 507 | } |
| 508 | *prefix_out = MakeShape(dims); |
| 509 | for (int i = rank; i < rank_s; ++i) dims.push_back(Dim(s, i)); |
| 510 | *s_out = MakeShape(dims); |
| 511 | return Status::OK(); |
| 512 | } |
| 513 | |
| 514 | void InferenceContext::Relax(ShapeHandle s_old, ShapeHandle s_new, |
| 515 | ShapeHandle* out) { |