| 366 | } |
| 367 | |
| 368 | Status InferenceContext::WithRank(ShapeHandle shape, int64 rank, |
| 369 | ShapeHandle* out) { |
| 370 | if (rank > kint32max) { |
| 371 | return errors::InvalidArgument("Rank cannot exceed kint32max"); |
| 372 | } |
| 373 | const int32 existing = Rank(shape); |
| 374 | if (existing == rank) { |
| 375 | *out = shape; |
| 376 | return Status::OK(); |
| 377 | } |
| 378 | if (existing == kUnknownRank) { |
| 379 | std::vector<DimensionHandle> dims; |
| 380 | dims.reserve(rank); |
| 381 | for (int i = 0; i < rank; ++i) { |
| 382 | dims.push_back(UnknownDim()); |
| 383 | } |
| 384 | ShapeHandle shp = shape_manager_.MakeShape(dims); |
| 385 | return Merge(shape, shp, out); |
| 386 | } |
| 387 | *out = nullptr; |
| 388 | |
| 389 | return errors::InvalidArgument("Shape must be rank ", rank, " but is rank ", |
| 390 | existing); |
| 391 | } |
| 392 | |
| 393 | Status InferenceContext::WithRankAtLeast(ShapeHandle shape, int64 rank, |
| 394 | ShapeHandle* out) { |