From f81438cb0272b2fd45d83e56999e12182484ea4b Mon Sep 17 00:00:00 2001
From: Roman Dzhabarov <rdzhabarov@google.com>
Date: Wed, 27 Jan 2021 10:52:04 -0800
Subject: [PATCH] [MlirGraphOptimizationPass] fix the indexing issue into the
 per_pass_state.

PiperOrigin-RevId: 354125177
Change-Id: Ifcf01771e9c44d330a2abe5059bfec8aa55cadd7
---
 tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
index b8b7b0da1aa..3b713fc0140 100644
--- a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
+++ b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
@@ -207,8 +207,7 @@ Status MlirFunctionOptimizationPass::Run(
          overall_state == MlirOptimizationPassState::ShadowEnabled)) {
       pass_status =
           pass_registration.pass->Run(config_proto, *module_ref, **graph);
-    } else if (per_pass_state[per_pass_state_index] ==
-               MlirOptimizationPassState::ShadowEnabled) {
+    } else if (pass_state == MlirOptimizationPassState::ShadowEnabled) {
       // Make sure that the pass does not modify MLIR module if it's shadow
       // enabled.
       auto module_ref_clone = module_ref->clone();