} : (tensor<2x7xi64>) -> (tensor<2x7xi32>)
  %1 = "tf.Identity"(%0) : (tensor<2x7xi32>) -> tensor<2x7xi32>
  func.return %1 : tensor<2x7xi32>
}

// CHECK-LABEL: func @func
// CHECK-SAME: %arg0: tensor<2x7xi64> {mhlo.sharding = "\08\03\1A\03\01\01\02\22\02\00\01B\01\00"
func.func @func(%arg0: tensor<2x7xi64>) -> (tensor<2x7xi32>) {

    return %0#0, %0#1, %0#2, %0#3, %0#4 : tensor<?xi32>, tensor<?x!tf_type.string>, tensor<?x2xi32>, tensor<3x!tf_type.string>, tensor<?x3xi32>
  }

  // CHECK-LABEL: func @deplicated_return_from_host_and_tpu

  %group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
  // CHECK: "mhlo.all_reduce"
  // CHECK{LITERAL}: replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
  // CHECK-NOT: channel_handle
  // CHECK: mhlo.add
  %0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
  func.return %0 : tensor<f32>
}

// -----

  %0 = stablehlo.constant dense<2.000000e+00> : tensor<2x3xf32>
  %1 = stablehlo.constant dense<2.000000e+00> : tensor<1x3xf32>
  %2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
  %3 = stablehlo.add %2, %1 : tensor<1x3xf32>
  func.return %3: tensor<1x3xf32>
}

// CHECK-LABEL: func private @reorder_execute_arg_defining_ops
// CHECK: tf.VarHandleOp
// CHECK-NEXT: tf.ReadVariableOp
// CHECK-NEXT: tf.TPUCompileMlirAndExecute
func.func private @reorder_execute_arg_defining_ops(%arg0: tensor<1x3xf32> {tf.device = "/CPU:0"}) -> (tensor<1x1xf32> {tf.device = "/TPU:0"}) {

func.func @main(%arg0: tensor<5x7xf32>) -> tensor<5x7xf32> {
  func.return %arg0: tensor<5x7xf32>
// CHECK-LABEL: main
// CHECK: return %arg0 : tensor<5x7xf32>
}

// - transpose
//
func.func @transpose_2d(%arg0: tensor<2x3xf32>) -> tensor<3x2xf32> {
  %0 = "mhlo.transpose"(%arg0) <{permutation = dense<[1, 0]> : tensor<2xi64>}> : (tensor<2x3xf32>) -> tensor<3x2xf32>
  func.return %0 : tensor<3x2xf32>

    %rhs_dilation = "tf.Const"() {value = dense<1> : tensor<3xi32>} : () -> tensor<3xi32>
    %lhs_dilation = "tf.Const"() {value = dense<[4, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32>
    %padding = "tf.Const"() {value = dense<0> : tensor<3x2xi32>} : () -> tensor<3x2xi32>
    %strides = "tf.Const"() {value = dense<[3, 1, 1]> : tensor<3xi32>} : () -> tensor<3xi32>

func.func @select_float(%arg0: tensor<1x3xi1>, %arg1: tensor<1x3xf32>, %arg2: tensor<1x3xf32>) -> tensor<1x3xf32> {
  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<1x3xi1>, tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
  return %0 : tensor<1x3xf32>
}
// CHECK-LABEL: select_float
// CHECK-NOT: tfl.select_v2
// CHECK: stablehlo.select

// -----

}

// CHECK-LABEL: QuantizeSVDF
func.func @QuantizeSVDF(%arg0: tensor<1x3xf32>) -> tensor<1x2xf32>  {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[2.07937503, 1.365000e+01]> : tensor<2xf32>} : (tensor<1x3xf32>) -> tensor<1x3xf32>
  %1 = "tfl.pseudo_const"() {value = dense<[[1.125947117805481, 1.0, 1.1], [-1.164743185043335, -1.0, -1.1]]> : tensor<2x3xf32>} : () -> tensor<2x3xf32>

  func.return %2 : tensor<1x128x32xf32>

// CHECK-DAG:  [[VAL_1:%.*]] = arith.constant dense<0> : tensor<3xi32>
// CHECK-DAG:  [[VAL_2:%.*]] = arith.constant dense<[1, 128, 32]> : tensor<3xi32>
// CHECK:  [[VAL_3:%.*]] = "tfl.slice"(%arg0, [[VAL_1]], [[VAL_2]]) : (tensor<4x128x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x128x32xf32>
}

// -----

// CHECK-LABEL: @WhileCanonicalizeBug