ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf.contrib.checkpoint.split_dependency

View source on GitHub

Creates multiple dependencies with a synchronized save/restore.

Useful when a single op produces Tensors which should each be saved under different objects, or when Tensors saved with many different objects need to be restored together as inputs to a single op (i.e. an object which uses a single fused op may be swapped out for a subgraph of objects, and these two programs are checkpoint compatible).

component_names A sequence of names for the split dependencies. fill_save_buffer_fn must add these keys to the dictionary it is passed, and consume_restore_buffer_fn will receive a dictionary with these keys.
component_dtypes Data types for the Tensors being saved and restored, a sequence corresponding to component_names.
fill_save_buffer_fn A function which takes an empty dictionary as an argument and adds Tensors with component_names as keys. These Tensors will be saved as if they were individual variables.
consume_restore_buffer_fn A function which takes a dictionary with component_names as keys mapping to restored individual Tensors and returns a restore op (or if executing eagerly, runs the restoration and may return None).
device The device on which to run save and restore operations.

A dictionary mapping from names to Trackable objects. If one is reachable from an object as a dependency, the others should be too; adding dependencies on some but not all of the objects will result in errors.