| 58 | } |
| 59 | |
| 60 | bool TensorSliceSet::QueryMeta( |
| 61 | const TensorSlice& slice, |
| 62 | std::vector<std::pair<TensorSlice, string>>* results) const { |
| 63 | results->clear(); |
| 64 | Status s; |
| 65 | string str = slice.DebugString(); |
| 66 | // First we check if there is an exactly match (this is the dominant case). |
| 67 | const TensorSliceSet::SliceInfo* info = gtl::FindOrNull(slices_, str); |
| 68 | if (info) { |
| 69 | results->emplace_back(std::make_pair(info->slice, info->tag)); |
| 70 | return true; |
| 71 | } else { |
| 72 | // We didn't find any exact match but there is still a possibility that |
| 73 | // multiple existing slices can be patched together to output the slice. |
| 74 | // We figure this out by computing the intersection of each of the existing |
| 75 | // slices with the query slice, and check if the union of all these |
| 76 | // intersections cover the entire slice. We rely on the fact that the |
| 77 | // existing slices don't have any intersection among themselves. |
| 78 | TensorShape target_shape; |
| 79 | Status s; |
| 80 | s = slice.SliceTensorShape(shape_, &target_shape); |
| 81 | if (!s.ok()) { |
| 82 | LOG(WARNING) << s; |
| 83 | return false; |
| 84 | } |
| 85 | int64 total_size = target_shape.num_elements(); |
| 86 | |
| 87 | int64 overlap_size = 0; |
| 88 | TensorSlice intersection; |
| 89 | TensorShape inter_shape; |
| 90 | for (const auto& x : slices_) { |
| 91 | if (slice.Intersect(x.second.slice, &intersection)) { |
| 92 | s = intersection.SliceTensorShape(shape_, &inter_shape); |
| 93 | if (!s.ok()) { |
| 94 | LOG(WARNING) << s; |
| 95 | return false; |
| 96 | } |
| 97 | overlap_size += inter_shape.num_elements(); |
| 98 | results->emplace_back(std::make_pair(x.second.slice, x.second.tag)); |
| 99 | } |
| 100 | } |
| 101 | if (total_size == overlap_size) { |
| 102 | // We have it! |
| 103 | return true; |
| 104 | } else { |
| 105 | // We don't have all the data for the asked tensor slice |
| 106 | results->clear(); |
| 107 | return false; |
| 108 | } |
| 109 | } |
| 110 | } |
| 111 | |
| 112 | Status RegisterTensorSlice( |
| 113 | const string& name, const TensorShape& shape, DataType type, |