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[tensor_lib.Tensor, tf_variables.Variable]]
)

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

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