SaveableObject implementation that handles BloomFilterAdmitStrategy
| 71 | |
| 72 | |
| 73 | class BloomFilterSaveable(saveable_object.SaveableObject): |
| 74 | """SaveableObject implementation that handles BloomFilterAdmitStrategy""" |
| 75 | custom_restore = True |
| 76 | |
| 77 | def __init__(self, bloom_filter): |
| 78 | self._init = bloom_filter.false_initializer |
| 79 | self._handle = bloom_filter.handle |
| 80 | slice_spec = "{} {},{}".format(bloom_filter.max_slice_size, |
| 81 | bloom_filter.slice_offset, bloom_filter.slice_size) |
| 82 | self._restore_name = bloom_filter.distributed_name |
| 83 | spec = saveable_object.SaveSpec(self._handle, slice_spec, self._restore_name) |
| 84 | self._restore_slice = slice_spec |
| 85 | super(BloomFilterSaveable, self).__init__(self._handle, [spec], bloom_filter.name) |
| 86 | |
| 87 | def restore_specs(self): |
| 88 | return [] |
| 89 | |
| 90 | def restore(self, restored_tensors, restored_shapes): |
| 91 | with ops.colocate_with(self._handle): |
| 92 | with ops.control_dependencies([self._init]): |
| 93 | return gen_io_ops.restore_bloom_filter( |
| 94 | restored_tensors[0], self._restore_name, self._restore_slice, |
| 95 | self._handle) |
| 96 | |
| 97 | class HashTableSaveable(saveable_object.SaveableObject): |
| 98 | """SaveableObject implementation that handles HashTable.""" |