diff --git a/tensorflow/compiler/mlir/lite/tests/analyze-variables.mlir b/tensorflow/compiler/mlir/lite/tests/analyze-variables.mlir
index ff0856ffdeb56d..fdf4af46e55452 100644
--- a/tensorflow/compiler/mlir/lite/tests/analyze-variables.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/analyze-variables.mlir
@@ -57,7 +57,7 @@ module {
 
 // -----
 
-// CHECK: module attributes {tfl._legalize_tfl_variables = false}
+// CHECK: module attributes {tfl._legalize_tfl_variables = true}
 module {
   func @main() -> tensor<i32> {
     %0 = "tf.VarHandleOp"() {container = "c", shared_name = "v"} : () -> tensor<*x!tf_type.resource<tensor<*xi32>>>
diff --git a/tensorflow/compiler/mlir/lite/transforms/analyze_variables.cc b/tensorflow/compiler/mlir/lite/transforms/analyze_variables.cc
index 0d40878317f8ad..56ee3f377821e6 100644
--- a/tensorflow/compiler/mlir/lite/transforms/analyze_variables.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/analyze_variables.cc
@@ -33,9 +33,11 @@ bool IsSupportedTFLiteResourceOp(Operation* op) {
                    TF::LookupTableSizeV2Op>(op);
 }
 
-// Returns true if 'op' is a TFLite control flow operation.
-bool IsTFLiteControlFlowOp(Operation* op) {
-  return llvm::isa<TFL::WhileOp, TFL::IfOp, TFL::CallOnceOp>(op);
+// Returns true if 'op' is TF/TFLite control flow op that can accept resource
+// type. Usually these ops are just pass through, they call another subgraph and
+// pass the operands to.
+bool IsSupportedTFLiteControlFlow(Operation* op) {
+  return llvm::isa<TFL::WhileOp, TFL::IfOp, TF::IfOp, TFL::CallOnceOp>(op);
 }
 }  // namespace
 
@@ -67,15 +69,10 @@ class AnalyzeVariablesPass
     module.walk([&](Operation* op) {
       // Skip ops that are supported natively by TFLite.
       if (IsSupportedTFLiteResourceOp(op)) return WalkResult::advance();
+      if (IsSupportedTFLiteControlFlow(op)) return WalkResult::advance();
 
       // Check for ops that are legalized to TFLite.
       if (op->getDialect()->getNamespace() == "tfl") {
-        // TODO(b/189370197): Enable control flow ops after updating
-        // checks to handle them.
-        if (IsTFLiteControlFlowOp(op)) {
-          legalize_to_tfl = false;
-          return WalkResult::interrupt();
-        }
         return WalkResult::advance();
       }
       // Check for ops that are not legalized to TFLite.