1# Copyright 2022 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""DTensor Python API.""" 16 17from tensorflow.dtensor.python import gen_dtensor_ops as ops 18from tensorflow.dtensor.python import mesh_util 19from tensorflow.dtensor.python import tpu_util 20 21from tensorflow.dtensor.python.api import call_with_layout 22from tensorflow.dtensor.python.api import check_layout 23from tensorflow.dtensor.python.api import client_id 24from tensorflow.dtensor.python.api import copy_to_mesh 25from tensorflow.dtensor.python.api import device_name 26from tensorflow.dtensor.python.api import fetch_layout 27from tensorflow.dtensor.python.api import full_job_name 28from tensorflow.dtensor.python.api import heartbeat_enabled 29from tensorflow.dtensor.python.api import job_name 30from tensorflow.dtensor.python.api import jobs 31from tensorflow.dtensor.python.api import local_devices 32from tensorflow.dtensor.python.api import num_clients 33from tensorflow.dtensor.python.api import num_global_devices 34from tensorflow.dtensor.python.api import num_local_devices 35from tensorflow.dtensor.python.api import pack 36from tensorflow.dtensor.python.api import relayout 37from tensorflow.dtensor.python.api import run_on 38from tensorflow.dtensor.python.api import unpack 39from tensorflow.dtensor.python.d_checkpoint import DTensorCheckpoint 40from tensorflow.dtensor.python.d_variable import DVariable 41from tensorflow.dtensor.python.input_util import DTensorDataset 42from tensorflow.dtensor.python.input_util import TFDataServiceConfig 43from tensorflow.dtensor.python.layout import Layout 44from tensorflow.dtensor.python.layout import MATCH 45from tensorflow.dtensor.python.layout import Mesh 46from tensorflow.dtensor.python.layout import UNSHARDED 47from tensorflow.dtensor.python.save_restore import enable_save_as_bf16 48from tensorflow.dtensor.python.save_restore import name_based_restore 49from tensorflow.dtensor.python.save_restore import name_based_save 50from tensorflow.dtensor.python.save_restore import sharded_save 51