tf.experimental.dtensor.sharded_save

Saves given named tensor slices in a sharded, multi-client safe fashion.

tf.experimental.dtensor.sharded_save(
    mesh: tf.experimental.dtensor.Mesh,
    file_prefix: Union[str, tf.Tensor],
    tensor_names: Union[List[str], tf.Tensor],
    shape_and_slices: Union[List[str], tf.Tensor],
    tensors: List[Union[tensor_lib.Tensor, tf_variables.Variable]]
)

The method makes sure the checkpoint directory state is correct in a sharded mutli-client saving. Namely, we place a barrier after SaveV2 to make sure every client has done writing the files. And another one after MergeV2Checkpoints to make sure all Metadata is properly merged.

Upon existing, the checkpoint is completed and the all directory operations are done.

meshThe Mesh that contains the Tensors to save.
file_prefixThe prefix of checkpoint.
tensor_namesa list of tensor names used in save op.
shape_and_slicesa list of shape and slice specification used in save op. The only supported value is "" as we don't support distributed saving with slices yet.
tensorsa list of tensors used in save op. The order should match tensor_names.

A MergeV2Checkpoints op that merged all Metadata.