Add func namespace in MLIR files

This will eventually be required rather than having an implicit namespace.

PiperOrigin-RevId: 436843763
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/basic_lstm.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/basic_lstm.mlir
index 7310502..9a8fbf7 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/basic_lstm.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/basic_lstm.mlir
@@ -6,5 +6,5 @@
 // CHECK: "tfl.basic_lstm"({{.*}}) {asymmetric_quantize_inputs = false, cell_clip = 1.000000e+00 : f32, fused_activation_function = "RELU", kernel_type = "BASIC", proj_clip = 2.000000e+00 : f32} : (tensor<1x384xf32>, tensor<1x96xf32>, tensor<384x480xf32>, tensor<384xf32>, tensor<1x96xf32>) -> (tensor<1x96xf32>, tensor<1x96xf32>, tensor<1x480xf32>, tensor<1x384xf32>)
 
   %0:4 = "tfl.basic_lstm"(%arg0, %arg1, %arg2, %arg3, %arg4) {fused_activation_function = "RELU", cell_clip = 1.0 : f32, proj_clip = 2.0 : f32} : (tensor<1x384xf32>, tensor<1x96xf32>, tensor<384x480xf32>, tensor<384xf32>, tensor<1x96xf32>) -> (tensor<1x96xf32>, tensor<1x96xf32>, tensor<1x480xf32>, tensor<1x384xf32>)
-  return %0#0 : tensor<1x96xf32>
+  func.return %0#0 : tensor<1x96xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/bucketize.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/bucketize.mlir
index b065039..bb3f582 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/bucketize.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/bucketize.mlir
@@ -5,5 +5,5 @@
   // CHECK-LABEL: @main
   // CHECK: "tfl.bucketize"(%arg0) {boundaries = [0.000000e+00 : f32, 1.000000e+01 : f32, 1.000000e+02 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32>
   %0 = "tfl.bucketize"(%arg0) {boundaries = [0.0 : f32, 10.0 : f32, 100.0 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32>
-  return %0 : tensor<3x2xi32>
+  func.return %0 : tensor<3x2xi32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/constants.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/constants.mlir
index ce53a4b..9832a09 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/constants.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/constants.mlir
@@ -5,56 +5,56 @@
   // CHECK-LABEL: @bool
   // CHECK: value = dense<[false, true, true, false]> : tensor<4xi1>
   %0 = "tfl.pseudo_const"() { value = dense<[false, true, true, false]> : tensor<4xi1> } : () -> tensor<4xi1>
-  return %0 : tensor<4xi1>
+  func.return %0 : tensor<4xi1>
 }
 
 func @complex64() -> tensor<4xcomplex<f32>> {
   // CHECK-LABEL: @complex64
   // CHECK: value = opaque<"tf", "0x746674656E736F722464747970653A2044545F434F4D504C455836342074656E736F725F7368617065207B2064696D207B2073697A653A2034207D207D2074656E736F725F636F6E74656E743A20225C3030305C3030305C3230303F5C3030305C3030305C3230303F5C3030305C3030305C303030405C3030305C3030305C303030405C3030305C30303040405C3030305C30303040405C3030305C3030305C323030405C3030305C3030305C3230304022"> : tensor<4xcomplex<f32>>
   %0 = "tfl.pseudo_const"() { value = opaque<"tf", "0x746674656E736F722464747970653A2044545F434F4D504C455836342074656E736F725F7368617065207B2064696D207B2073697A653A2034207D207D2074656E736F725F636F6E74656E743A20225C3030305C3030305C3230303F5C3030305C3030305C3230303F5C3030305C3030305C303030405C3030305C3030305C303030405C3030305C30303040405C3030305C30303040405C3030305C3030305C323030405C3030305C3030305C3230304022"> : tensor<4xcomplex<f32>> } : () -> tensor<4xcomplex<f32>>
-  return %0 : tensor<4xcomplex<f32>>
+  func.return %0 : tensor<4xcomplex<f32>>
 }
 
 func @complex128() -> tensor<4xcomplex<f64>> {
   // CHECK-LABEL: @complex128
   // CHECK: value = opaque<"tf", "0x746674656E736F722464747970653A2044545F434F4D504C45583132382074656E736F725F7368617065207B2064696D207B2073697A653A2034207D207D2074656E736F725F636F6E74656E743A20225C3030305C3030305C3030305C3030305C3030305C3030305C3336303F5C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303130405C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303230405C3030305C3030305C3030305C3030305C3030305C3030305C3030304022"> : tensor<4xcomplex<f64>>
   %0 = "tfl.pseudo_const"() { value = opaque<"tf", "0x746674656E736F722464747970653A2044545F434F4D504C45583132382074656E736F725F7368617065207B2064696D207B2073697A653A2034207D207D2074656E736F725F636F6E74656E743A20225C3030305C3030305C3030305C3030305C3030305C3030305C3336303F5C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303130405C3030305C3030305C3030305C3030305C3030305C3030305C303030405C3030305C3030305C3030305C3030305C3030305C3030305C303230405C3030305C3030305C3030305C3030305C3030305C3030305C3030304022"> : tensor<4xcomplex<f64>> } : () -> tensor<4xcomplex<f64>>
-  return %0 : tensor<4xcomplex<f64>>
+  func.return %0 : tensor<4xcomplex<f64>>
 }
 
 func @f16() -> tensor<4xf16> {
   // CHECK-LABEL: @f16
   // CHECK: value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<4xf16>
   %0 = "tfl.pseudo_const"() { value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf16> } : () -> tensor<4xf16>
-  return %0 : tensor<4xf16>
+  func.return %0 : tensor<4xf16>
 }
 
 func @f32() -> tensor<4xf32> {
   // CHECK-LABEL: @f32
   // CHECK: value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<4xf32>
   %0 = "tfl.pseudo_const"() { value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32> } : () -> tensor<4xf32>
-  return %0 : tensor<4xf32>
+  func.return %0 : tensor<4xf32>
 }
 
 func @f64() -> tensor<4xf64> {
   // CHECK-LABEL: @f64
   // CHECK: value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<4xf64>
   %0 = "tfl.pseudo_const"() { value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf64> } : () -> tensor<4xf64>
-  return %0 : tensor<4xf64>
+  func.return %0 : tensor<4xf64>
 }
 
 func @i8() -> tensor<4xi8> {
   // CHECK-LABEL: @i8
   // CHECK: value = dense<[1, 2, 3, 4]> : tensor<4xi8>
   %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 4]> : tensor<4xi8> } : () -> tensor<4xi8>
-  return %0 : tensor<4xi8>
+  func.return %0 : tensor<4xi8>
 }
 
 func @i16() -> tensor<4xi16> {
   // CHECK-LABEL: @i16
   // CHECK: value = dense<[1, 2, 3, 258]> : tensor<4xi16>
   %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 258]> : tensor<4xi16> } : () -> tensor<4xi16>
-  return %0 : tensor<4xi16>
+  func.return %0 : tensor<4xi16>
 }
 
 func @i32() -> tensor<4xi32> {
@@ -62,80 +62,80 @@
   // CHECK: value = dense<[1, 2, 3, 16909060]> : tensor<4xi32>
   // Check bytes come back in the right order
   %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 16909060]> : tensor<4xi32> } : () -> tensor<4xi32>
-  return %0 : tensor<4xi32>
+  func.return %0 : tensor<4xi32>
 }
 
 func @i64() -> tensor<4xi64> {
   // CHECK-LABEL: @i64
   // CHECK: value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64>
   %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 72623859790382856]> : tensor<4xi64> } : () -> tensor<4xi64>
-  return %0 : tensor<4xi64>
+  func.return %0 : tensor<4xi64>
 }
 
 func @string() -> tensor<2x2x!tf_type.string> {
   // CHECK-LABEL: @string
   // CHECK: value = dense<{{\[\["1", "12"\], \["123", "1234"\]\]}}> : tensor<2x2x!tf_type.string>
   %0 = "tfl.pseudo_const"() { value = dense<[["1", "12"], ["123", "1234"]]> : tensor<2x2x!tf_type.string> } : () -> tensor<2x2x!tf_type.string>
-  return %0 : tensor<2x2x!tf_type.string>
+  func.return %0 : tensor<2x2x!tf_type.string>
 }
 
 func @string_norank() -> tensor<!tf_type.string> {
   // CHECK-LABEL: @string_norank
   // CHECK: value = dense<"test"> : tensor<!tf_type.string>
   %0 = "tfl.pseudo_const"() { value = dense<"test"> : tensor<!tf_type.string> } : () -> tensor<!tf_type.string>
-  return %0 : tensor<!tf_type.string>
+  func.return %0 : tensor<!tf_type.string>
 }
 
 func @uint8() -> tensor<4xui8> {
   // CHECK-LABEL: @uint8
   // CHECK: value = dense<[222, 173, 190, 239]> : tensor<4xui8>
   %0 = "tfl.pseudo_const"() {value = dense<[222, 173, 190, 239]> : tensor<4xui8>} : () -> tensor<4xui8>
-  return %0 : tensor<4xui8>
+  func.return %0 : tensor<4xui8>
 }
 
 func @qi32_per_axis() -> tensor<3x3x!quant.uniform<i32:f32:1, {1.0, 0.5:1, 0.25:1}>> {
   // CHECK-LABEL: @qi32_per_axis
   // CHECK: {qtype = tensor<3x3x!quant.uniform<i32:f32:1, {1.000000e+00,5.000000e-01:1,2.500000e-01:1}>>, value = dense<1> : tensor<3x3xi32>} : () -> tensor<3x3x!quant.uniform<i32:f32:1, {1.000000e+00,5.000000e-01:1,2.500000e-01:1}>>
   %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x3x!quant.uniform<i32:f32:1, {1.0, 0.5:1, 0.25:1}>>, value = dense<1> : tensor<3x3xi32>} : () -> tensor<3x3x!quant.uniform<i32:f32:1, {1.0, 0.5:1, 0.25:1}>>
-  return %0 : tensor<3x3x!quant.uniform<i32:f32:1, {1.0, 0.5:1, 0.25:1}>>
+  func.return %0 : tensor<3x3x!quant.uniform<i32:f32:1, {1.0, 0.5:1, 0.25:1}>>
 }
 
 func @qi32_per_axis_zero() -> tensor<3x3x!quant.uniform<i32:f32:0, {1.0, 0.5:1, 0.25:1}>> {
   // CHECK-LABEL: @qi32_per_axis_zero
   // CHECK: {qtype = tensor<3x3x!quant.uniform<i32:f32:0, {1.000000e+00,5.000000e-01:1,2.500000e-01:1}>>, value = dense<1> : tensor<3x3xi32>} : () -> tensor<3x3x!quant.uniform<i32:f32:0, {1.000000e+00,5.000000e-01:1,2.500000e-01:1}>>
   %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x3x!quant.uniform<i32:f32:0, {1.0, 0.5:1, 0.25:1}>>, value = dense<1> : tensor<3x3xi32>} : () -> tensor<3x3x!quant.uniform<i32:f32:0, {1.0, 0.5:1, 0.25:1}>>
-  return %0 : tensor<3x3x!quant.uniform<i32:f32:0, {1.0, 0.5:1, 0.25:1}>>
+  func.return %0 : tensor<3x3x!quant.uniform<i32:f32:0, {1.0, 0.5:1, 0.25:1}>>
 }
 
 func @qu8() -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>> {
   // CHECK-LABEL: @qu8
   // CHECK: {qtype = tensor<3x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>, value = dense<1> : tensor<3xi8>} : () -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>
   %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>, value = dense<1> : tensor<3xi8>} : () -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
-  return %0 : tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
+  func.return %0 : tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
 }
 
 func @sparse_f32() -> tensor<3x2xf32> {
   // CHECK-LABEL: @sparse_f32
   // CHECK: {compressed_data = dense<[1.000000e+00, 2.000000e+00, 5.000000e-01, 2.500000e-01, -1.000000e+00, -2.000000e+00, -5.000000e-01, -2.500000e-01]> : tensor<8xf32>, s_param = {block_map = [3 : i32, 1 : i32], dim_metadata = [{dense_size = 16 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 0 : i32, format = "SPARSE_CSR", indices = [1 : i32, 4 : i32, 9 : i32], segments = [0 : i32, 5 : i32, 11 : i32]}], traversal_order = [0 : i32, 1 : i32, 2 : i32, 3 : i32]}, value = dense<0.000000e+00> : tensor<3x2xf32>}
   %0 = "tfl.pseudo_sparse_const"() {compressed_data = dense<[1.0, 2.0, 0.5, 0.25, -1.0, -2.0, -0.5, -0.25]> : tensor<8xf32>, s_param = {block_map = [3 : i32, 1 : i32], dim_metadata = [{dense_size = 16 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 0 : i32, format = "SPARSE_CSR", indices = [1 : i32, 4 : i32, 9 : i32], segments = [0 : i32, 5 : i32, 11 : i32]}], traversal_order = [0 : i32, 1 : i32, 2 : i32, 3 : i32]}, value = dense<0.000000e+00> : tensor<3x2xf32>} : () -> tensor<3x2xf32>
-  return %0: tensor<3x2xf32>
+  func.return %0: tensor<3x2xf32>
 }
 
 func @sparse_f16() -> tensor<3x2xf16> {
   // CHECK-LABEL: @sparse_f16
   // CHECK: {compressed_data = dense<[1.000000e+00, 2.000000e+00, 5.000000e-01, 2.500000e-01, -1.000000e+00, -2.000000e+00, -5.000000e-01, -2.500000e-01]> : tensor<8xf16>, s_param = {block_map = [3 : i32, 1 : i32], dim_metadata = [{dense_size = 16 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 0 : i32, format = "SPARSE_CSR", indices = [1 : i32, 4 : i32, 9 : i32], segments = [0 : i32, 5 : i32, 11 : i32]}], traversal_order = [0 : i32, 1 : i32, 2 : i32, 3 : i32]}, value = dense<0.000000e+00> : tensor<3x2xf16>}
   %0 = "tfl.pseudo_sparse_const"() {compressed_data = dense<[1.0, 2.0, 0.5, 0.25, -1.0, -2.0, -0.5, -0.25]> : tensor<8xf16>, s_param = {block_map = [3 : i32, 1 : i32], dim_metadata = [{dense_size = 16 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 0 : i32, format = "SPARSE_CSR", indices = [1 : i32, 4 : i32, 9 : i32], segments = [0 : i32, 5 : i32, 11 : i32]}], traversal_order = [0 : i32, 1 : i32, 2 : i32, 3 : i32]}, value = dense<0.000000e+00> : tensor<3x2xf16>} : () -> tensor<3x2xf16>
-  return %0: tensor<3x2xf16>
+  func.return %0: tensor<3x2xf16>
 }
 
 func @sparse_qu8() -> tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>> {
   // CHECK-LABEL: @sparse_qu8
   // CHECK: {compressed_data = dense<[1, 2, 3, 4, -1, -2, -3, -4]> : tensor<8xi8>, qtype = tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>, s_param = {block_map = [3 : i32, 1 : i32], dim_metadata = [{dense_size = 16 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 0 : i32, format = "SPARSE_CSR", indices = [1 : i32, 4 : i32, 9 : i32], segments = [0 : i32, 5 : i32, 11 : i32]}], traversal_order = [0 : i32, 1 : i32, 2 : i32, 3 : i32]}, value = dense<0> : tensor<3x2xi8>}
   %0 = "tfl.pseudo_sparse_qconst"() {compressed_data = dense<[1, 2, 3, 4, -1, -2, -3, -4]> : tensor<8xi8>, qtype = tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>, s_param = {block_map = [3 : i32, 1 : i32], dim_metadata = [{dense_size = 16 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 1 : i32, format = "DENSE", indices = [], segments = []}, {dense_size = 0 : i32, format = "SPARSE_CSR", indices = [1 : i32, 4 : i32, 9 : i32], segments = [0 : i32, 5 : i32, 11 : i32]}], traversal_order = [0 : i32, 1 : i32, 2 : i32, 3 : i32]}, value = dense<42> : tensor<3x2xi8>} : () -> tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>
-  return %0: tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>
+  func.return %0: tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>>
 }
 
 // Identity function to make the exporter happy
 func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> {
-  return %arg0 : tensor<4xi8>
+  func.return %arg0 : tensor<4xi8>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/custom_op.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/custom_op.mlir
index 0e0b21f..e674113 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/custom_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/custom_op.mlir
@@ -2,7 +2,7 @@
 
 func @main(%arg0: tensor<32x4x4x128xf32>, %arg1: tensor<1x32x42x128xf32>, %arg2: tensor<4xi32>) -> tensor<1x64x84x32xf32> {
   %0 = "tfl.custom"(%arg0, %arg1, %arg2) {custom_code = "Convolution2DTransposeBias", custom_option = opaque<"tfl", "0x010000000200000002000000"> : tensor<12xi8>} : (tensor<32x4x4x128xf32>, tensor<1x32x42x128xf32>, tensor<4xi32>) -> tensor<1x64x84x32xf32>
-  return %0 : tensor<1x64x84x32xf32>
+  func.return %0 : tensor<1x64x84x32xf32>
 }
 // CHECK-LABEL: main
 // CHECK: "tfl.custom"(%arg0, %arg1, %arg2) {custom_code = "Convolution2DTransposeBias", custom_option = opaque<"tfl", "0x010000000200000002000000"> : tensor<12xi8>} : (tensor<32x4x4x128xf32>, tensor<1x32x42x128xf32>, tensor<4xi32>) -> tensor<1x64x84x32xf32>
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dynamic_shape.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dynamic_shape.mlir
index 2e79dd4..1502464 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dynamic_shape.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/dynamic_shape.mlir
@@ -5,5 +5,5 @@
   %cst = arith.constant dense<1.0> : tensor<4xf32>
   %cst_3 = arith.constant dense<2.0> : tensor<4x3x3x3xf32>
   %0 = "tfl.conv_2d"(%arg0, %cst_3, %cst) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU6", padding = "VALID", stride_h = 2 : i32, stride_w = 2 : i32} : (tensor<?x19x19x3xf32>, tensor<4x3x3x3xf32>, tensor<4xf32>) -> tensor<?x9x9x4xf32>
-  return %0 : tensor<?x9x9x4xf32>
+  func.return %0 : tensor<?x9x9x4xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/external_constant.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/external_constant.mlir
index 5ce8f1a..8e236b9 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/external_constant.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/external_constant.mlir
@@ -5,7 +5,7 @@
 ^bb0(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>):
   %cst = arith.constant dense<1.0> : tensor<40xf32>
   %0:2 = "tfl.fully_connected"(%arg0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, tensor<40xf32>) -> (tensor<40x40xf32>, tensor<40x40xf32>)
-  return %0 : tensor<40x40xf32>
+  func.return %0 : tensor<40x40xf32>
 
 // CHECK-LABEL: func @main(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>) -> tensor<40x40xf32>
 // CHECK:      %[[CONST:[0-9]+]] = "tfl.external_const"() {buffer_index = 3 : i32} : () -> tensor<40xf32>
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir
index c064bb8..548fa2d 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/if_op.mlir
@@ -4,15 +4,15 @@
 // CHECK:   %{{.*}} = "tf.If"(%{{.*}}, %{{.*}}, %{{.*}}) {else_branch = @cond_false, is_stateless = false, then_branch = @cond_true} : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
   %0 = "tfl.less"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>
   %1 = "tf.If"(%0, %arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = false} : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
-  return %1 : tensor<1xf32>
+  func.return %1 : tensor<1xf32>
 }
 
 func @cond_true(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
   %0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<*xf32>
-  return %0 : tensor<*xf32>
+  func.return %0 : tensor<*xf32>
 }
 
 func @cond_false(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
   %0 = tfl.mul %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<*xf32>
-  return %0 : tensor<*xf32>
+  func.return %0 : tensor<*xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/importer_test_min_max.cc.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/importer_test_min_max.cc.mlir
index b95b409..eb7616b 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/importer_test_min_max.cc.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/importer_test_min_max.cc.mlir
@@ -5,7 +5,7 @@
 ^bb0(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>):
   %cst = arith.constant dense<1.0> : tensor<40xf32>
   %0:2 = "tfl.fully_connected"(%arg0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, tensor<40xf32>) -> (tensor<40x40xf32>, tensor<40x40xf32>)
-  return %0 : tensor<40x40xf32>
+  func.return %0 : tensor<40x40xf32>
 
 // CHECK-LABEL: func @main(%arg0: tensor<40x37xf32>, %arg1: tensor<40x37xf32>) -> tensor<40x40xf32>
 // CHECK:      %[[CONST:[0-9]+]] = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<40xf32>}
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_arrays.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_arrays.mlir
index 4581caf..bd7e8ea 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_arrays.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_arrays.mlir
@@ -5,7 +5,7 @@
   %0 = "tfl.pseudo_const" () {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> loc("Const")
   %1 = "tfl.squared_difference"(%arg0, %0) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("squared_difference")
   %2 = "tfl.mul"(%0, %1) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("mul")
-  return %2 : tensor<4xf32>
+  func.return %2 : tensor<4xf32>
 
 // CHECK-LABEL: main
 // CHECK-NOT: tfl.squared_difference
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_output_names_attr.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_output_names_attr.mlir
index 81b9a31..8929440 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_output_names_attr.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/input_output_names_attr.mlir
@@ -8,5 +8,5 @@
 attributes {tf.entry_function = {inputs = "input0,input1", outputs = "output0,output1"}} {
   %0 = "tfl.neg"(%arg0) : (tensor<4xi8>) -> tensor<4xi8> loc("neg")
   %1 = "tfl.neg"(%arg1) : (tensor<4xi32>) -> tensor<4xi32> loc("neg")
-  return %1, %0 : tensor<4xi32>, tensor<4xi8>
+  func.return %1, %0 : tensor<4xi32>, tensor<4xi8>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/lstm.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/lstm.mlir
index 7f79f2f..639c77d 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/lstm.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/lstm.mlir
@@ -5,7 +5,7 @@
   %cst0 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %cst1 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %24 = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %cst0, %cst1, %arg18, %arg19, %arg20, %arg21) ({}) {cell_clip = 0.000000e+00 : f32, fused_activation_function = "NONE", kernel_type = "FULL", proj_clip = 0.000000e+00 : f32} : (tensor<1x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<1x4xf32>
-  return %24 : tensor<1x4xf32>
+  func.return %24 : tensor<1x4xf32>
 // CHECK-LABEL: main
 // seperate lines since there is no region for this op. third_party/tensorflow/compiler/mlir/lite/ir/tfl_ops.td: 3252
 // CHECK-DAG: %[[RES0:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
@@ -21,7 +21,7 @@
 func @testFullyQuantizedLSTM(%arg0: tensor<1x528x!quant.uniform<i8:f32, 0.037248000502586365:-19>>, %arg1: tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.059801999479532242>>, %arg2: tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.031925998628139496>>, %arg3: tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.056272000074386597>>, %arg4: tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.063763998448848724>>, %arg5: tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.013358999975025654>>, %arg6: tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.022830000147223473>>, %arg7: tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.032276000827550888>>, %arg8: tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.035427000373601913>>, %arg9: tensor<2048x!quant.uniform<i32:f32, 4.2675782196965883E-7>>, %arg10: tensor<2048x!quant.uniform<i32:f32, 1.0742187583900886E-7>>, %arg11: tensor<2048x!quant.uniform<i32:f32, 1.6406249869760359E-7>>, %arg12: tensor<2048x!quant.uniform<i32:f32, 1.523437447303877E-7>>, %arg13: tensor<640x2048x!quant.uniform<i8<-127:127>:f32, 0.021174000576138496>>, %arg14: tensor<640x!quant.uniform<i32:f32, 1.601389680352559E-4>>, %arg15: tensor<2048x!quant.uniform<i16:f32, 4.3700000969693065E-4>>, %arg16: tensor<2048x!quant.uniform<i16:f32, 1.1000000085914508E-4>>, %arg17: tensor<2048x!quant.uniform<i16:f32, 1.6799999866634607E-4>>, %arg18: tensor<2048x!quant.uniform<i16:f32, 1.55999994603917E-4>>, %arg19: tensor<1x640x!quant.uniform<i8:f32, 0.09671100229024887:10>>, %arg20: tensor<1x2048x!quant.uniform<i16:f32, 4.8799999058246613E-4>>) -> tensor<1x640x!quant.uniform<i8:f32, 0.09671100229024887:10>> {
     %cst = "tfl.no_value"() {value = unit} : () -> none
     %0 = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %cst, %cst, %cst, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg19, %arg20, %arg15, %arg16, %arg17, %arg18) ({}) {cell_clip = 1.000000e+01 : f32, fused_activation_function = "TANH", input_to_input_intermediate = tensor<0x!quant.uniform<i16:f32, 0.0049890000373125076>>, input_to_forget_intermediate = tensor<0x!quant.uniform<i16:f32, 0.0078849997371435165>>, input_to_cell_intermediate = tensor<0x!quant.uniform<i16:f32, 0.0087630003690719604>>, input_to_output_intermediate = tensor<0x!quant.uniform<i16:f32, 0.0057529998011887074>>, effective_hidden_scale_intermediate = tensor<0x!quant.uniform<i8:f32, 0.0075630000792443752:2>>, kernel_type = "FULL", proj_clip = 0.01 : f32} : (tensor<1x528x!quant.uniform<i8:f32, 0.037248000502586365:-19>>, tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.059801999479532242>>, tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.031925998628139496>>, tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.056272000074386597>>, tensor<2048x528x!quant.uniform<i8<-127:127>:f32, 0.063763998448848724>>, tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.013358999975025654>>, tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.022830000147223473>>, tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.032276000827550888>>, tensor<2048x640x!quant.uniform<i8<-127:127>:f32, 0.035427000373601913>>, none, none, none, tensor<2048x!quant.uniform<i32:f32, 4.2675782196965883E-7>>, tensor<2048x!quant.uniform<i32:f32, 1.0742187583900886E-7>>, tensor<2048x!quant.uniform<i32:f32, 1.6406249869760359E-7>>, tensor<2048x!quant.uniform<i32:f32, 1.523437447303877E-7>>, tensor<640x2048x!quant.uniform<i8<-127:127>:f32, 0.021174000576138496>>, tensor<640x!quant.uniform<i32:f32, 1.601389680352559E-4>>, tensor<1x640x!quant.uniform<i8:f32, 0.09671100229024887:10>>, tensor<1x2048x!quant.uniform<i16:f32, 4.8799999058246613E-4>>, tensor<2048x!quant.uniform<i16:f32, 4.3700000969693065E-4>>, tensor<2048x!quant.uniform<i16:f32, 1.1000000085914508E-4>>, tensor<2048x!quant.uniform<i16:f32, 1.6799999866634607E-4>>, tensor<2048x!quant.uniform<i16:f32, 1.55999994603917E-4>>) -> tensor<1x640x!quant.uniform<i8:f32, 0.09671100229024887:10>>
-    return %0 : tensor<1x640x!quant.uniform<i8:f32, 0.09671100229024887:10>>
+    func.return %0 : tensor<1x640x!quant.uniform<i8:f32, 0.09671100229024887:10>>
 // CHECK-LABEL: testFullyQuantizedLSTM
 // CHECK: %[[CST:.*]] = "tfl.no_value"() {value} : () -> none
 // CHECK: %[[RES0:.*]] = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %[[CST]], %[[CST]], %[[CST]], %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg19, %arg20, %arg15, %arg16, %arg17, %arg18)
@@ -34,7 +34,7 @@
 func @testUnidirectionalSequenceLstmWithIntermediates(%arg0: tensor<? x ? x ? x f32>, %arg1: tensor<? x ? x f32>, %arg2: tensor<? x ? x f32>, %arg3: tensor<? x ? x f32>, %arg4: tensor<? x ? x f32>, %arg5: tensor<? x ? x f32>, %arg6: tensor<? x ? x f32>, %arg7: tensor<? x ? x f32>, %arg8: tensor<? x ? x f32>, %arg9: tensor<? x f32>, %arg10: tensor<? x f32>, %arg11: tensor<? x f32>, %arg12: tensor<? x f32>, %arg13: tensor<? x f32>, %arg14: tensor<? x f32>, %arg15: tensor<? x f32>, %arg16: tensor<? x ? x f32>, %arg17: tensor<? x f32>, %arg18: tensor<? x ? x f32>, %arg19: tensor<? x ? xf32>, %arg20: tensor<? x f32>, %arg21: tensor<? x f32>, %arg22: tensor<? x f32>, %arg23: tensor<? x f32>) -> tensor<? x ? x ? x f32> {
   // CHECK: "tfl.unidirectional_sequence_lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23) {asymmetric_quantize_inputs = false, cell_clip = 1.000000e+01 : f32, effective_hidden_scale_intermediate = tensor<0x!quant.uniform<i8:f32, 0.0077881771139800549>>, fused_activation_function = "TANH", input_to_cell_intermediate = tensor<0xf32>, input_to_forget_intermediate = tensor<0xf32>, input_to_input_intermediate = tensor<0xf32>, input_to_output_intermediate = tensor<0xf32>, proj_clip = 0.000000e+00 : f32, time_major = false} : (tensor<?x?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> tensor<?x?x?xf32>
   %0 = "tfl.unidirectional_sequence_lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23) {cell_clip = 1.000000e+01 : f32, effective_hidden_scale_intermediate = tensor<0x!quant.uniform<i8:f32, 0.0077881771139800549>>, fused_activation_function = "TANH", input_to_cell_intermediate = tensor<0xf32>, input_to_forget_intermediate = tensor<0xf32>, input_to_input_intermediate = tensor<0xf32>, input_to_output_intermediate = tensor<0xf32>, proj_clip = 0.000000e+00 : f32, time_major = false} : (tensor<?x?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) -> tensor<?x?x?xf32>
-  return %0 : tensor<?x?x?xf32>
+  func.return %0 : tensor<?x?x?xf32>
 }
 
 // -----
@@ -44,7 +44,7 @@
   %cst0 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %cst1 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %24 = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %cst0, %cst1, %arg18, %arg19, %arg20, %arg21) ({}) {asymmetric_quantize_inputs = true, cell_clip = 0.000000e+00 : f32, fused_activation_function = "NONE", kernel_type = "FULL", proj_clip = 0.000000e+00 : f32} : (tensor<1x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<1x4xf32>
-  return %24 : tensor<1x4xf32>
+  func.return %24 : tensor<1x4xf32>
 
 // CHECK-DAG: %[[RES0:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
 // CHECK-DAG: %[[RES1:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
@@ -61,7 +61,7 @@
   %cst0 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %cst1 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %24 = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %cst0, %cst1, %arg18, %arg19, %arg20, %arg21) ({}) {asymmetric_quantize_inputs = false, cell_clip = 0.000000e+00 : f32, fused_activation_function = "NONE", kernel_type = "FULL", proj_clip = 0.000000e+00 : f32} : (tensor<1x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<1x4xf32>
-  return %24 : tensor<1x4xf32>
+  func.return %24 : tensor<1x4xf32>
 
 // CHECK-DAG: %[[RES0:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
 // CHECK-DAG: %[[RES1:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
@@ -78,7 +78,7 @@
   %cst0 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %cst1 = "tfl.pseudo_const" () {value = dense<0.0> : tensor<1x4xf32>} : () -> tensor<1x4xf32> loc("Const")
   %24 = "tfl.lstm"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7, %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15, %arg16, %arg17, %cst0, %cst1, %arg18, %arg19, %arg20, %arg21) ({}) {cell_clip = 0.000000e+00 : f32, fused_activation_function = "NONE", kernel_type = "FULL", proj_clip = 0.000000e+00 : f32} : (tensor<1x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<1x4xf32>, tensor<1x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<1x4xf32>
-  return %24 : tensor<1x4xf32>
+  func.return %24 : tensor<1x4xf32>
 
 // CHECK-DAG: %[[RES0:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
 // CHECK-DAG: %[[RES1:.*]] = "tfl.pseudo_const"() {value = dense<{{.*}}> : tensor<1x4xf32>} : () -> tensor<1x4xf32>
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/many_attribute_op.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/many_attribute_op.mlir
index e0de33e..e94227f 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/many_attribute_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/many_attribute_op.mlir
@@ -5,5 +5,5 @@
 ^bb0(%arg0: tensor<1x6x6x16xf32>):
   // CHECK: "tfl.average_pool_2d"(%{{.*}}) {filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32} : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32>
   %0 = "tfl.average_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32} : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32> loc("avgpool")
-  return %0 : tensor<1x1x1x16xf32>
+  func.return %0 : tensor<1x1x1x16xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/math.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/math.mlir
index fe39d46..999d400 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/math.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/math.mlir
@@ -14,5 +14,5 @@
   %3 = "tfl.div"(%2, %1) {fused_activation_function = "NONE"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32> loc("div")
   %4 = "tfl.exp"(%3) : (tensor<4xf32>) -> tensor<4xf32> loc("exp")
   %5 = "tfl.neg"(%4) : (tensor<4xf32>) -> tensor<4xf32> loc("neg")
-  return %5 : tensor<4xf32>
+  func.return %5 : tensor<4xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/matmul.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/matmul.mlir
index 53f6b9f..458e5cca 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/matmul.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/matmul.mlir
@@ -2,7 +2,7 @@
 
 func @main(%arg0: tensor<4x10x15xf32>, %arg1: tensor<4x15x17xf32>) -> tensor<4x10x17xf32> {
   %0 = "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
-  return %0:  tensor<4x10x17xf32>
+  func.return %0:  tensor<4x10x17xf32>
 
 // CHECK-LABEL: main
 // CHECK: %[[RESULT0:.*]] =  "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
@@ -12,7 +12,7 @@
 // CHECK-LABEL: testMatmulAsymAttributeTrue
 func @testMatmulAsymAttributeTrue(%arg0: tensor<4x10x15xf32>, %arg1: tensor<4x15x17xf32>) -> tensor<4x10x17xf32> {
   %0 = "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = true} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
-  return %0:  tensor<4x10x17xf32>
+  func.return %0:  tensor<4x10x17xf32>
 
 // CHECK: %[[RESULT0:.*]] =  "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = true} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
 // CHECK: return %[[RESULT0]]
@@ -21,7 +21,7 @@
 // CHECK-LABEL: testMatmulAsymAttributeFalse
 func @testMatmulAsymAttributeFalse(%arg0: tensor<4x10x15xf32>, %arg1: tensor<4x15x17xf32>) -> tensor<4x10x17xf32> {
   %0 = "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
-  return %0:  tensor<4x10x17xf32>
+  func.return %0:  tensor<4x10x17xf32>
 
 // CHECK: %[[RESULT0:.*]] =  "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
 // CHECK: return %[[RESULT0]]
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/optional.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/optional.mlir
index a13858c..32b83f2 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/optional.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/optional.mlir
@@ -7,5 +7,5 @@
 // CHECK-SAME: (tensor<40x37xf32>, tensor<40x37xf32>, none) -> (tensor<40x40xf32>, tensor<40x40xf32>)
   %cst = "tfl.no_value"() {value = unit} : () -> none
   %0:2 = "tfl.fully_connected"(%arg0, %arg1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<40x37xf32>, tensor<40x37xf32>, none) -> (tensor<40x40xf32>, tensor<40x40xf32>)
-  return %0 : tensor<40x40xf32>
+  func.return %0 : tensor<40x40xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/output_arrays.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/output_arrays.mlir
index ba97d86..ee90ac0 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/output_arrays.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/output_arrays.mlir
@@ -15,5 +15,5 @@
   // CHECK: %[[NEG:.*]] = "tfl.neg"
   %5 = "tfl.neg"(%4) : (tensor<4xf32>) -> tensor<4xf32> loc("neg")
   // CHECK: return %[[MUL]], %[[EXP]], %[[DIV]]
-  return %5 : tensor<4xf32>
+  func.return %5 : tensor<4xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning.mlir
index 9ceda83..262b7e6 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning.mlir
@@ -15,5 +15,5 @@
   // CHECK-NOT: "tfl.neg"
   %5 = "tfl.neg"(%4) : (tensor<4xf32>) -> tensor<4xf32> loc("neg")
   // CHECK: return %[[MUL]], %[[EXP]], %[[DIV]]
-  return %5 : tensor<4xf32>
+  func.return %5 : tensor<4xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning_function_input_as_output.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning_function_input_as_output.mlir
index 18fac00..91d6843 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning_function_input_as_output.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/pruning_function_input_as_output.mlir
@@ -12,5 +12,5 @@
   // CHECK-NOT: "tfl.neg"
   %4 = "tfl.neg"(%3) : (tensor<4xf32>) -> tensor<4xf32> loc("neg")
   // CHECK: return %[[ARG]], %[[DIV]], %[[EXP]]
-  return %4 : tensor<4xf32>
+  func.return %4 : tensor<4xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quant_stats.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quant_stats.mlir
index 3f73e44..68f4ef8 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quant_stats.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quant_stats.mlir
@@ -6,5 +6,5 @@
 // CHECK: %[[RES0:.*]] = "quant.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32>
 
   %0 = "quant.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32>
-  return %0 : tensor<1x512x672x8xf32>
+  func.return %0 : tensor<1x512x672x8xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir
index 22c9d7c..9098c1d 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/quantization.mlir
@@ -16,7 +16,7 @@
   %4 = "tfl.reshape"(%3, %cst) : (tensor<1x112x112x32x!quant.uniform<u8:f32, 0.023528476789885875>>, tensor<2xi32>) -> tensor<1x401408x!quant.uniform<u8:f32, 0.023528476789885875>>
   %5 = "tfl.softmax"(%4) {beta = 1.000000e+00 : f32} : (tensor<1x401408x!quant.uniform<u8:f32, 0.023528476789885875>>) -> tensor<1x401408x!quant.uniform<u8:f32, 3.906250e-03>>
   %6 = "tfl.dequantize"(%5) : (tensor<1x401408x!quant.uniform<u8:f32, 3.906250e-03>>) -> tensor<1x401408xf32>
-  return %6 : tensor<1x401408xf32>
+  func.return %6 : tensor<1x401408xf32>
 }
 
 // CHECK-LABEL: quantized_constant
@@ -25,7 +25,7 @@
   %cst = "tfl.pseudo_qconst"() {qtype = tensor<1x2x!quant.uniform<u8:f32, 1.0>>, value = dense<-76> : tensor<1x2xi8>} : () -> tensor<1x2x!quant.uniform<u8:f32, 1.0>>
   %2 = "tfl.concatenation"(%1, %cst) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2x!quant.uniform<u8:f32, 1.0>>, tensor<1x2x!quant.uniform<u8:f32, 1.0>>) -> tensor<2x2x!quant.uniform<u8:f32, 1.0>>
   %3 = "tfl.dequantize"(%2) : (tensor<2x2x!quant.uniform<u8:f32, 1.0>>) -> tensor<2x2xf32>
-  return %3 : tensor<2x2xf32>
+  func.return %3 : tensor<2x2xf32>
 
 // CHECK-NEXT: %[[Q:.*]] = "tfl.quantize"(%arg0) {qtype = tensor<1x2x!quant.uniform<u8:f32, 1.000000e+00>>} : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<u8:f32, 1.000000e+00>>
 // CHECK-NEXT: %[[CST:.*]] = "tfl.pseudo_qconst"() {qtype = tensor<1x2x!quant.uniform<u8:f32, 1.000000e+00>>, value = dense<-76> : tensor<1x2xi8>} : () -> tensor<1x2x!quant.uniform<u8:f32, 1.000000e+00>>
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/reshape.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/reshape.mlir
index 12e0fab..5f76cd4 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/reshape.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/reshape.mlir
@@ -7,5 +7,5 @@
   %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
   %0 = "tfl.pseudo_const" () {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> loc("Const")
   %1 = "tfl.reshape" (%0, %cst) : (tensor<4xf32>, tensor<2xi32>) -> tensor<2x2xf32> loc("reshape")
-  return %1 : tensor<2x2xf32>
+  func.return %1 : tensor<2x2xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature.mlir
index bd1ae49..f7d1768 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature.mlir
@@ -13,6 +13,6 @@
     %cst_1 = arith.constant dense<1.0> : tensor<5x384xf32>
     %0 = "tfl.fully_connected"(%arg0, %cst_0, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<?x384xf32>, tensor<5x384xf32>, tensor<5xf32>) -> tensor<?x5xf32>
     %1 = "tfl.fully_connected"(%arg0, %cst_1, %cst) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<?x384xf32>, tensor<5x384xf32>, tensor<5xf32>) -> tensor<?x5xf32>
-    return %1, %0 : tensor<?x5xf32>, tensor<?x5xf32>
+    func.return %1, %0 : tensor<?x5xf32>, tensor<?x5xf32>
   }
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature_with_multiple_entry_points.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature_with_multiple_entry_points.mlir
index 5e1c289..9a5e859 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature_with_multiple_entry_points.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/signature_with_multiple_entry_points.mlir
@@ -18,11 +18,11 @@
   func @add(%arg0: tensor<?xf32> {tf_saved_model.index_path = ["input1"]}, %arg1: tensor<?xf32> {tf_saved_model.index_path = ["input2"]}) -> (tensor<?xf32> {tf_saved_model.index_path = ["result"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "input1:0,input2:0", outputs = "result:0"}, tf_saved_model.exported_names = ["add"]} {
     %0 = tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<?xf32>
     %1 = tfl.add %0, %0 {fused_activation_function = "NONE"} : tensor<?xf32>
-    return %1 : tensor<?xf32>
+    func.return %1 : tensor<?xf32>
   }
 
   func @sub(%arg0: tensor<?xf32> {tf_saved_model.index_path = ["input2"]}, %arg1: tensor<?xf32> {tf_saved_model.index_path = ["input1"]}) -> (tensor<?xf32> {tf_saved_model.index_path = ["result"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "input2:0,input1:0", outputs = "result:0"}, tf_saved_model.exported_names = ["sub"]} {
     %0 = tfl.sub %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<?xf32>
-    return %0 : tensor<?xf32>
+    func.return %0 : tensor<?xf32>
   }
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/simple.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/simple.mlir
index 53ffc63..54101b6 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/simple.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/simple.mlir
@@ -17,5 +17,5 @@
   %1 = "tfl.sub" (%arg0, %0) {fused_activation_function = "RELU6"} : (tensor<3x2xi32>, tensor<3x2xi32>) -> tensor<3x2xi32> loc("sub")
   %2 = "arith.constant" () {value = dense<10> : tensor<i32>} : () -> tensor<i32> loc("Const2")
   %3 = "tfl.add" (%2, %1) {fused_activation_function = "NONE"} : (tensor<i32>, tensor<3x2xi32>) -> tensor<3x2xi32> loc("add")
-  return %3 : tensor<3x2xi32>
+  func.return %3 : tensor<3x2xi32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/unranked_function_output.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/unranked_function_output.mlir
index 1a0c889..0209f91 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/unranked_function_output.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/unranked_function_output.mlir
@@ -16,13 +16,13 @@
 // CHECK: }
 func @main(%arg0: tensor<1xf32>) -> tensor<*xf32> {
   %0 = "tf.While"(%arg0) {cond = @cond, body = @body, is_stateless = false} : (tensor<1xf32>) -> tensor<*xf32>
-  return %0 : tensor<*xf32>
+  func.return %0 : tensor<*xf32>
 }
 
 func @cond(%arg1: tensor<*xf32>) -> tensor<*xf32> {
-  return %arg1: tensor<*xf32>
+  func.return %arg1: tensor<*xf32>
 }
 
 func @body(%arg1: tensor<*xf32>) -> tensor<*xf32> {
-  return %arg1: tensor<*xf32>
+  func.return %arg1: tensor<*xf32>
 }
diff --git a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/while_op.mlir b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/while_op.mlir
index 6d1effc..220aa06 100644
--- a/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/while_op.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/flatbuffer2mlir/while_op.mlir
@@ -20,18 +20,18 @@
     %1:2 = call @body(%arg2, %arg3) : (tensor<*xi32>, tensor<*xf32>) -> (tensor<*xi32>, tensor<*xf32>)
     "tfl.yield"(%1#0, %1#1) : (tensor<*xi32>, tensor<*xf32>) -> ()
   }) {is_stateless = false} : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)
-  return %0#1 : tensor<1xf32>
+  func.return %0#1 : tensor<1xf32>
 }
 
 func @cond(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>) -> tensor<i1> {
   %cst = arith.constant dense<0> : tensor<i32> loc("Const")
   %0 = "tfl.greater"(%arg0, %cst) : (tensor<*xi32>, tensor<i32>) -> tensor<i1>
-  return %0 : tensor<i1>
+  func.return %0 : tensor<i1>
 }
 
 func @body(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>) -> (tensor<*xi32>, tensor<*xf32>) {
   %cst = arith.constant dense<1> : tensor<i32> loc("Const")
   %0 = "tfl.sub"(%arg0, %cst) {fused_activation_function = "NONE"} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
   %1 = tfl.add %arg1, %arg1 {fused_activation_function = "NONE"} : tensor<*xf32>
-  return %0, %1 : tensor<*xi32>, tensor<*xf32>
+  func.return %0, %1 : tensor<*xi32>, tensor<*xf32>
 }