Queries the cache. Parameters ---------- keys : Tensor The keys to query the cache with. offset : int The offset to be added to the keys. Default is 0. Returns ------- tuple(Tensor, Tensor, Tensor, Tensor)
(self, keys, offset=0)
| 69 | return self._cache.nbytes |
| 70 | |
| 71 | def query(self, keys, offset=0): |
| 72 | """Queries the cache. |
| 73 | |
| 74 | Parameters |
| 75 | ---------- |
| 76 | keys : Tensor |
| 77 | The keys to query the cache with. |
| 78 | offset : int |
| 79 | The offset to be added to the keys. Default is 0. |
| 80 | |
| 81 | Returns |
| 82 | ------- |
| 83 | tuple(Tensor, Tensor, Tensor, Tensor) |
| 84 | A tuple containing |
| 85 | (values, missing_indices, missing_keys, missing_offsets) where |
| 86 | values[missing_indices] corresponds to cache misses that should be |
| 87 | filled by quering another source with missing_keys. If keys is |
| 88 | pinned, then the returned values tensor is pinned as well. The |
| 89 | missing_offsets tensor has the partition offsets of missing_keys. |
| 90 | """ |
| 91 | self.total_queries += keys.shape[0] |
| 92 | ( |
| 93 | positions, |
| 94 | index, |
| 95 | missing_keys, |
| 96 | found_pointers, |
| 97 | found_offsets, |
| 98 | missing_offsets, |
| 99 | ) = self._policy.query(keys, offset) |
| 100 | values = self._cache.query(positions, index, keys.shape[0]) |
| 101 | self._policy.reading_completed(found_pointers, found_offsets) |
| 102 | self.total_miss += missing_keys.shape[0] |
| 103 | missing_index = index[positions.size(0) :] |
| 104 | return values, missing_index, missing_keys, missing_offsets |
| 105 | |
| 106 | def query_and_replace(self, keys, reader_fn, offset=0): |
| 107 | """Queries the cache. Then inserts the keys that are not found by |