mesh_util_test.py
unittest TestCase cleanup contextPython unittest import
SHA-256ebce91304bd684e07055d97ddfcef0999bd67b652977d8c4922692de96fd41fc
MaleculeH₂(Cm₂F)Md₂(In₃Pa)
Evidence
1# Copyright 2023 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 …
11:46… D, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Tests for mesh_util."""
16import os
17from unittest import mock
18
19from absl.testing import parameterized
20
21from tensorflow.dtensor.python import accelerator_util
22from tensorflow.dtensor.python import config
23from tensorflow.dtensor.python import mesh_util
24from tensorflow.dtensor.python.tests import test_util
25from tensorflow.python.eager import context
26from tensorflow.python.framework import config as tf_config
27from tensorflow.python.framework import device as tf_device
28from tensorflow.python.platform import test
29
30
31class MeshUtilTest(test_util.DTensorBaseTest):
32 """Tests for mesh_util that do not require accelerator initialization."""
33
34 def test_mesh_creation(self):
35 self.skipForDeviceType(
36 ['TPU'], reason='Test is intended for CPUs and G …
62:19… ('do_not_use_xla_spmd', False))
63 def test_tpu_2d_mesh_creation(self, use_xla_spmd):
64 self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended fo …
218:51… ration()] * 4
219 )
220 with mock.patch.object(config, 'client_id', return_value=0):
221 mesh_1 = mesh_util.create_distributed_mesh(
222 mesh_name='multi_client_1d_mesh_1',
223 …
242:8… elerator_util.initialize_accelerator_system(device_type)
243
244 def tearDown(self):
245 super().tearDown()
246 context._reset_context() # pylint: disable=protected-access
247
248 def test_is_initialized(self) …