func.func private @if_then_branch(%arg0: tensor<*xi1>, %arg1: tensor<*xi64>, %arg2: tensor<*xi64>) -> tensor<*xi1> {
    func.return %arg0 : tensor<*xi1>
  }
  func.func private @if_else_branch(%arg0: tensor<*xi1>, %arg1: tensor<*xi64>, %arg2: tensor<*xi64>) -> tensor<*xi1> {
    func.return %arg0 : tensor<*xi1>
  }
  func.func @"tpu_subgraph"(%arg0: tensor<*xi64>) -> tensor<*xi1> {

// CHECK-LABEL: func @no_spmd
// CHECK-SAME: ([[ARG0:%.+]]: tensor<!tf_type.resource<tensor<i32>>>, [[ARG1:%.+]]: tensor<!tf_type.resource<tensor<i32>>>)
func.func @no_spmd(%arg0: tensor<!tf_type.resource<tensor<i32>>>, %arg1: tensor<!tf_type.resource<tensor<i32>>>) {
  // CHECK:      "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])

func.func @main(%arg0: tensor<32x4x4x128xf32>, %arg1: tensor<1x32x42x128xf32>, %arg2: tensor<4xi32>) -> tensor<1x64x84x32xf32> {

  func.func private @_tf_func(%arg0: tensor<?xi32>, %arg1: tensor<*xi32>) {
    // CHECK: tf.StreamResults

    // StreamResults is a pseudo op in this test.
    "tf.StreamResults"(%arg0, %arg1) : (tensor<?xi32>, tensor<*xi32>) -> ()
    func.return
  }

  // CHECK-LABEL: func @main
  // CHECK-SAME:    %[[ARG0:.*]]: tensor<10xi32>, %[[ARG1:.*]]: tensor<10xi32>

    func.return %1: tensor<*xf32>
  }
  func.func private @composite_matmul_fn_1(%arg0: tensor<2x12xf32>, %arg1: tensor<12x2xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {
    %0 = "tf.MatMul"(%arg0, %arg1) {attr_map = "0:transpose_a,1:transpose_b", device = "", transpose_a = false, transpose_b = false} : (tensor<2x12xf32>, tensor<12x2xf32>) -> tensor<*xf32>
    return %0 : tensor<*xf32>
  }
}

  %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>
  func.return %0, %1 : tensor<*xi32>, tensor<*xf32>
}

func.func @main(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0:2 = "tfl.while"(%arg0, %arg1) ({
  ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xf32>):

func.func @testDilatedConv(%arg0: tensor<1x128x128x3xf32>, %arg1: tensor<5x5x3x8xf32>) -> tensor<1x120x120x8xf32> {
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %cst_0 = arith.constant dense<4> : tensor<2x2xi32>
  %0 = "tf.SpaceToBatchND"(%arg0, %cst, %cst_0) : (tensor<1x128x128x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x68x68x3xf32>

  %2 = "mhlo.pad"(%arg0, %1) {edge_padding_high = dense<[0, 0, 1, 1]> : tensor<4xi64>, edge_padding_low = dense<[0, 0, 1, 1]> : tensor<4xi64>, interior_padding = dense<0> : tensor<4xi64>} : (tensor<1x3x6x6xf32>, tensor<f32>) -> tensor<1x3x8x8xf32>
  %3 = "mhlo.reduce_window"(%2, %1) ({
  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
    %7 = mhlo.add %arg1, %arg2 : tensor<f32>
    mhlo.return %7 : tensor<f32>

// TUPLE-ARGS:         %[[ARG_TUPLE]] = (f32[], f32[]) parameter(0)
// TUPLE-ARGS:         [[ARG0:%.*]] = f32[] get-tuple-element((f32[], f32[]) %[[ARG_TUPLE]]), index=0
// TUPLE-ARGS:         [[ARG1:%.*]] = f32[] get-tuple-element((f32[], f32[]) %[[ARG_TUPLE]]), index=1
// TUPLE-ARGS:         [[ADD:%.*]] = f32[] add(f32[] [[ARG0]], f32[] [[ARG1]])
// TUPLE-ARGS:         ROOT %tuple.{{[0-9]+}} = (f32[]) tuple(f32[] [[ADD]])

// -----

func.func @handles_regions_that_use_arguments(%arg0: f32, %arg1: f32) {
  %0 = "glue.someop"() : () -> f32
  "foo.someop"() ({
    "bar.a"(%arg0) : (f32) -> f32
    "bar.b"(%arg1) : (f32) -> f32
    "bar.c"(%0) : (f32) -> f32
  }, {}): () -> ()
  return
}

// CHECK: func @handles_regions_that_use_arguments
// CHECK: call [[foo:@[^(]*]](%arg0, %arg1, %0)
// CHECK: func [[bar:@[^(]*]]