1# Copyright 2017 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 16"""Helper library for functions used during TPU compilation.""" 17 18import contextlib 19import threading 20 21 22class TpuContext(threading.local): 23 """A context object holding state about the TPU computation being built.""" 24 25 def __init__(self): 26 """Creates a new TpuContext.""" 27 self._number_of_shards = None 28 29 @property 30 def number_of_shards(self): 31 return self._number_of_shards 32 33 def set_number_of_shards(self, number_of_shards): 34 self._number_of_shards = number_of_shards 35 36 37# The Tpu context holds the number of shards when a sharded computation is 38# being built, or None if no computation is being built. 39_current_tpu_context = TpuContext() 40 41 42@contextlib.contextmanager 43def tpu_shard_context(number_of_shards): 44 """A context manager setting current number of shards.""" 45 if _current_tpu_context.number_of_shards is not None: 46 raise NotImplementedError( 47 "tpu_shard_context cannot be nested." 48 "If you're using TPUEstimator with inference_on_tpu, " 49 "make sure you have set " 50 "export_saved_model_api_version=ExportSavedModelApiVersion.V2 in " 51 "the creation of TPUEstimator.") 52 try: 53 _current_tpu_context.set_number_of_shards(number_of_shards) 54 yield 55 finally: 56 _current_tpu_context.set_number_of_shards(None) 57 58 59def get_tpu_context(): 60 return _current_tpu_context 61 62 63# Decorator function for tpu computation func that was passed to tpu.rewrite() 64# if there is an embedded training loop in this func, trace tools will generate 65# step markers for each iteration. 66def on_device_training_loop(func): 67 # Value for this attribute is from xla.DebugOptions.StepMarkerLocation. 68 setattr(func, "step_marker_location", "STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP") 69 return func 70