tf.experimental.dtensor.run_on

Выполняет заключённые функции в области устройства DTensor.

@contextlib.contextmanager
tf.experimental.dtensor.run_on(
    layout_or_mesh: Union[tf.experimental.dtensor.Layout, tf.experimental.dtensor.Mesh]
)

Эта функция возвращает область. Все операции и tf.functions в этой области будут выполняться на устройстве DTensor с использованием сетки, предоставленной или прикреплённой к макету. Это полезно для обертывания любого tf.function, который не принимает DTensor в качестве входных данных, но хотел бы получить DTensor в качестве результата. Область также гарантирует, что все небольшие константы будут дублироваться как DTensor.