Allow the TPUConfigResourceMgr to be fetched without causing initialization.
PiperOrigin-RevId: 387673988
Change-Id: If0ae627182e014f075af1391cf3e8dfde9b9ac19
diff --git a/tensorflow/core/tpu/tpu_configuration.cc b/tensorflow/core/tpu/tpu_configuration.cc
index 3788d5c..750252f 100644
--- a/tensorflow/core/tpu/tpu_configuration.cc
+++ b/tensorflow/core/tpu/tpu_configuration.cc
@@ -31,8 +31,10 @@
void MaybeInitializeTPUSystemForTests() {}
#endif
-ResourceMgr* GetTPUConfigResourceMgr() {
- MaybeInitializeTPUSystemForTests();
+ResourceMgr* GetTPUConfigResourceMgr(bool initialize_first) {
+ if (initialize_first) {
+ MaybeInitializeTPUSystemForTests();
+ }
// Put all TPU-related state in the global ResourceMgr. This includes the
// TpuPodState, compilation cache, etc. We don't use the TPU_SYSTEM
diff --git a/tensorflow/core/tpu/tpu_configuration.h b/tensorflow/core/tpu/tpu_configuration.h
index 6c337bd..0fbdb0f 100644
--- a/tensorflow/core/tpu/tpu_configuration.h
+++ b/tensorflow/core/tpu/tpu_configuration.h
@@ -23,7 +23,7 @@
void MaybeInitializeTPUSystemForTests();
// Returns a process-wide global ResourceMgr.
-ResourceMgr* GetTPUConfigResourceMgr();
+ResourceMgr* GetTPUConfigResourceMgr(bool initialize_first = true);
} // namespace tensorflow
