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[ops.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.

mesh The Mesh that contains the Tensors to save.
file_prefix The prefix of checkpoint.
tensor_names a list of tensor names used in save op.
shape_and_slices a 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.
tensors a list of tensors used in save op. The order should match tensor_names.

A MergeV2Checkpoints op that merged all Metadata.