tf.experimental.dtensor.name_based_restore

Restores from checkpoint_prefix to name based DTensors.

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.