Allows export of functions capturing a Dataset in SavedModels. When saving a SavedModel, `tf.saved_model.save` traverses the object graph. Since Datasets reference _VariantTracker objects, that traversal will find a _VariantTracker for each Dataset and so know how to save and restore functi
| 3289 | |
| 3290 | |
| 3291 | class _VariantTracker(tracking.CapturableResource): |
| 3292 | """Allows export of functions capturing a Dataset in SavedModels. |
| 3293 | |
| 3294 | When saving a SavedModel, `tf.saved_model.save` traverses the object |
| 3295 | graph. Since Datasets reference _VariantTracker objects, that traversal will |
| 3296 | find a _VariantTracker for each Dataset and so know how to save and restore |
| 3297 | functions which reference the Dataset's variant Tensor. |
| 3298 | """ |
| 3299 | |
| 3300 | def __init__(self, variant_tensor, resource_creator): |
| 3301 | """Record that `variant_tensor` is associated with `resource_creator`. |
| 3302 | |
| 3303 | Args: |
| 3304 | variant_tensor: The variant-dtype Tensor associated with the Dataset. This |
| 3305 | Tensor will be a captured input to functions which use the Dataset, and |
| 3306 | is used by saving code to identify the corresponding _VariantTracker. |
| 3307 | resource_creator: A zero-argument function which creates a new |
| 3308 | variant-dtype Tensor. This function will be included in SavedModels and |
| 3309 | run to re-create the Dataset's variant Tensor on restore. |
| 3310 | """ |
| 3311 | super(_VariantTracker, self).__init__(device="CPU") |
| 3312 | self._resource_handle = variant_tensor |
| 3313 | self._create_resource = resource_creator |
| 3314 | |
| 3315 | |
| 3316 | def _is_padded_shape_compatible_with(padded_shape, input_component_shape): |