tf.experimental.dtensor.default_mesh

Задает сетку устройств DTensor по умолчанию для функций в данной области видимости.

@contextlib.contextmanager
tf.experimental.dtensor.default_mesh(
    mesh: tf.experimental.dtensor.Mesh
)

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