| 496 | |
| 497 | template <class Shape> |
| 498 | void TensorShapeBase<Shape>::set_dim(int d, int64 size) { |
| 499 | CHECK_GE(d, 0); |
| 500 | CHECK_LT(d, dims()); |
| 501 | CHECK_GE(size, 0); |
| 502 | if (tag() == REP16 && size < kMaxRep16) { |
| 503 | as16()->dims_[d] = |
| 504 | kIsPartial && size < 0 ? kUnknownRep16 : static_cast<uint16>(size); |
| 505 | } else if (tag() == REP32 && size < kMaxRep32) { |
| 506 | as32()->dims_[d] = |
| 507 | kIsPartial && size < 0 ? kUnknownRep32 : static_cast<uint32>(size); |
| 508 | } else if (tag() == REP_OUT_OF_LINE) { |
| 509 | (*as64()->dims_)[d] = size; |
| 510 | } else { |
| 511 | // Must upgrade |
| 512 | gtl::InlinedVector<int64, 8> vals; |
| 513 | AppendTo(*this, &vals); |
| 514 | vals[d] = size; |
| 515 | ClearAllButDataType(); |
| 516 | for (auto dval : vals) { |
| 517 | AddDim(dval); |
| 518 | } |
| 519 | } |
| 520 | TF_CHECK_OK(RecomputeNumElements()); |
| 521 | } |
| 522 | |
| 523 | template <class Shape> |
| 524 | Status TensorShapeBase<Shape>::SetDimWithStatus(int d, int64 size) { |