From 789ebcf00b3ceb8a4a7a073a271a89ed75f9734c Mon Sep 17 00:00:00 2001
From: Rick Chao <rchao@google.com>
Date: Thu, 28 May 2020 14:48:02 -0700
Subject: [PATCH] Add doc for extended._retrace_functions_for_each_device.
PiperOrigin-RevId: 313663754
Change-Id: Ie643c0e54adb3eed6a7d20ea3328f2d49fb8047c
---
tensorflow/python/distribute/distribute_lib.py | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/tensorflow/python/distribute/distribute_lib.py b/tensorflow/python/distribute/distribute_lib.py
index fbc8923e050..b77163cb97a 100644
--- a/tensorflow/python/distribute/distribute_lib.py
+++ b/tensorflow/python/distribute/distribute_lib.py
@@ -619,6 +619,11 @@ class StrategyBase(object):
if not hasattr(extended, "_retrace_functions_for_each_device"):
# pylint: disable=protected-access
+ # `extended._retrace_functions_for_each_device` dictates
+ # 1) whether all the ops created inside function will have devices
+ # inherited from outer stack, and
+ # 2) whether the same function will be retraced when it is called on
+ # different devices.
try:
extended._retrace_functions_for_each_device = (
len(extended.worker_devices) > 1)