| 110 | } |
| 111 | |
| 112 | Status RegisterTensorSlice( |
| 113 | const string& name, const TensorShape& shape, DataType type, |
| 114 | const string& tag, const TensorSlice& slice, |
| 115 | std::unordered_map<string, TensorSliceSet*>* tensor_slices) { |
| 116 | DCHECK_NE(tensor_slices, nullptr); |
| 117 | TensorSliceSet* tss = gtl::FindPtrOrNull(*tensor_slices, name); |
| 118 | // Create a tensor slice set if needed |
| 119 | if (!tss) { |
| 120 | tss = new TensorSliceSet(shape, type); |
| 121 | tensor_slices->insert(std::make_pair(name, tss)); |
| 122 | } else { |
| 123 | // Check if the shapes match |
| 124 | const TensorShape& tss_shape(tss->shape()); |
| 125 | if (!shape.IsSameSize(tss_shape)) { |
| 126 | return errors::Internal("Incompatible tensor shapes detected for tensor ", |
| 127 | name, ": existing = ", tss_shape.DebugString(), |
| 128 | ", new = ", shape.DebugString()); |
| 129 | } |
| 130 | if (type != tss->type()) { |
| 131 | return errors::Internal("Incompatible tensor types detected for tensor ", |
| 132 | name, |
| 133 | ": existing = ", DataTypeString(tss->type()), |
| 134 | ", new = ", DataTypeString(type)); |
| 135 | } |
| 136 | } |
| 137 | // Register the tensor slices without the actual data. |
| 138 | return tss->Register(slice, tag); |
| 139 | } |
| 140 | |
| 141 | } // namespace checkpoint |
| 142 |
no test coverage detected