tf.experimental.dtensor.name_based_save

Сохраняет тензоры с именами в контрольной точке.

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

Функция подготавливает входной словарь к формату sharded_save, чтобы использовать преимущества распределенного сохранения DTensor SPMD.

Как и при восстановлении, функция поддерживает сохранение только на одной сетке.