tf.experimental.dtensor.name_based_restore

Restores from checkpoint_prefix to name based DTensors.

tf.experimental.dtensor.name_based_restore(
    mesh: tf.experimental.dtensor.Mesh,
    checkpoint_prefix: str,
    name_tensor_dict: Dict[str, Union[ops.Tensor, tf_variables.Variable]]
)

It is required to have already-initialized DTensor variables that have same shape/dtype for the tensors being restored.

Also, we currently only support a named based restore on a single mesh.

mesh The single mesh that all Tensors would be restored to.
checkpoint_prefix The prefix of checkpoint to be restored.
name_tensor_dict A ordered dictionary of tensor_names to a DTensor. The DTensor shape/dtype must match the tensors being saved/restored for now.

A dictionary of name to its restored DTensor value.