%0 = "tfl.dequantize"(%arg0) : (tensor<1000x1000x!quant.uniform<i8:f32, 7.812500e-03>>) -> tensor<1000x1000xf32>
  %1 = "tfl.gather"(%0, %arg1) {axis = 0 : i32, batch_dims = 0 : i32}: (tensor<1000x1000xf32>, tensor<1x1xi32>) -> tensor<1x1x1000xf32>
  func.return %1 : tensor<1x1x1000xf32>

// CHECK-NEXT: tfl.gather
// CHECK-NEXT: tfl.dequantize
}

        %b = "tf.opB"(%a) : (tensor<i1>) -> tensor<i32>
        %c = "tf.opC"(%b) : (tensor<i32>) -> tensor<f32>
        tf_device.return %c, %b : tensor<f32>, tensor<i32>
      }) {device = "CPU:0"} : () -> (tensor<f32>, tensor<i32>)
      %d = "tf.opD"() : () -> tensor<i1>
      tf_executor.yield %a, %launch#0, %launch#1, %d : tensor<i1>, tensor<f32>, tensor<i32>, tensor<i1>
    }
    tf_executor.fetch
  }
  func.return
}

    %4:2 = tf_executor.island wraps "tf.Mul"(%2#1, %2#0) {T = "tfdtype$DT_INT32", device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32> loc("Square")
    %5:3 = tf_executor.Merge %3#0, %4#0 : tensor<i32> {device = "", N = 2, T = "tfdtype$DT_INT32"} loc("Merge")
    tf_executor.fetch
  }
  func.return

func.func @multiple_return(%arg0 : tensor<*xi32>, %arg1 : tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>) {
  %1 = "tf.Add"(%arg0, %arg1) {} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  %2 = "tf.Add"(%1, %arg1) {} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  func.return %1, %2 : tensor<*xi32>, tensor<*xi32>
}

// CHECK-LABEL: func @multiple_return
// CHECK-SAME: (%[[ARG0:.*]]: tensor<*xi32>, %[[ARG1:.*]]: tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>) {

      tf_executor.yield %1#1, %1#0 : tensor<i32>, tensor<i32>
    }
    tf_executor.fetch
  }
  func.return
}

// CKECK-SAME: <{calibration_method = 5 : i32, id = "skipping_id", num_bins = 32 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32}> : (tensor<1x3x4x3xf32>) -> (tensor<1x3x4x3xf32>, tensor<f32>, tensor<f32>, tensor<512xi64>)
// CHECK: %[[CUSTOM_AGGREGATOR_1:.*]], %[[MIN_1:.*]], %[[MAX_1:.*]], %[[HISTOGRAM_1:.*]] = "tf.CustomAggregator"

        %0 = "tf.Const"() {device = "", value = dense<0> : tensor<i32>} : () -> tensor<i32>
        "tf_executor.yield"(%0) : (tensor<i32>) -> ()
      }) : () -> (tensor<i32>, !tf_executor.control)
      %outputs_0:3, %control_1 = "tf_executor._SwitchN"(%outputs, %outputs) {T = i32, device = "", num_outs = 3 : i64} : (tensor<i32>, tensor<i32>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, !tf_executor.control)

// CHECK:           [[VAL_1:%.*]] = "tfl.average_pool_2d"([[VAL_0]]) <{filter_height = 48 : i32, filter_width = 48 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32>
// CHECK:           return [[VAL_1]] : tensor<1x1x1x512xf32>
// CHECK:         }

// -----

        %b = "tf.opB"(%a) : (tensor<i1>) -> tensor<i32>
        %c = "tf.opC"(%b) : (tensor<i32>) -> tensor<f32>
        tf_device.return %c, %b : tensor<f32>, tensor<i32>
      }) {device = "CPU:0"} : () -> (tensor<f32>, tensor<i32>)
      %d = "tf.opD"() : () -> tensor<i1>
      tf_executor.yield %a, %launch#0, %launch#1, %d : tensor<i1>, tensor<f32>, tensor<i32>, tensor<i1>
    }
    tf_executor.fetch
  }
  func.return
}

  %1:4 = "tf.CustomAggregator"(%arg0) {id = "1", calibration_method = 1 : i32, num_bins = 0 : i32, max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32} : (tensor<8x8x8x8xf32>) -> (tensor<8x8x8x8xf32>,...