@@ -33,16 +33,21 @@ void VLogGraphDebugString(Graph* g) {
3333
3434class EmbeddingForwardBackwardJointOptimizationPass : public GraphOptimizationPass {
3535 public:
36+ EmbeddingForwardBackwardJointOptimizationPass () : GraphOptimizationPass() {
37+ tensorflow::ReadBoolFromEnvVar (" TF_EMBEDDING_FBJ_OPT" ,
38+ /* default_val=*/ false , &embedding_fbj_opt_);
39+ if (!embedding_fbj_opt_) {
40+ VLOG (2 ) << " Graph Optimization Pass TF_EMBEDDING_FBJ_OPT is off." ;
41+ } else {
42+ VLOG (2 ) << " Graph Optimization Pass TF_EMBEDDING_FBJ_OPT is on." ;
43+ }
44+ }
45+
3646 Status Run (const GraphOptimizationPassOptions& options) override {
37- bool embedding_fbj_opt = false ;
38- TF_CHECK_OK (
39- tensorflow::ReadBoolFromEnvVar (" TF_EMBEDDING_FBJ_OPT" ,
40- /* default_val=*/ false , &embedding_fbj_opt));
41- if (!embedding_fbj_opt) {
42- LOG (INFO) << " TF_EMBEDDING_FBJ_OPT off." ;
47+ if (!embedding_fbj_opt_) {
4348 return Status::OK ();
4449 }
45- LOG (INFO) << " TF_EMBEDDING_FBJ_OPT on. " ;
50+
4651 if (options.graph == nullptr ) {
4752 // TODO(apassos) returning OK feels weird here as we can't do anything
4853 // without a graph, but some tests require this.
@@ -198,7 +203,10 @@ class EmbeddingForwardBackwardJointOptimizationPass : public GraphOptimizationPa
198203 return Status::OK ();
199204 }
200205
206+ private:
207+ bool embedding_fbj_opt_ = false ;
201208};
209+
202210REGISTER_OPTIMIZATION (OptimizationPassRegistry::POST_REWRITE_FOR_EXEC, 0 ,
203211 EmbeddingForwardBackwardJointOptimizationPass);
204212
0 commit comments