func.func @testTranspose(tensor<2x3xf32>, tensor<1x2xi32>) -> tensor<3x2xf32> {
^bb0(%arg0: tensor<2x3xf32>, %arg1: tensor<1x2xi32>):
  // expected-error @+1 {{expected perm to be a 1-D Tensor, got perm of rank 2}}
  %0 = "tf.Transpose"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", Tperm = "tfdtype$DT_INT32"} : (tensor<2x3xf32>, tensor<1x2xi32>) -> tensor<3x2xf32>
  func.return %0 : tensor<3x2xf32>
}

// -----

}

// -----

func.func @testConcat(%arg0: tensor<1x2xi32>, %arg1: tensor<1x2xi32>) -> tensor<2x2xi32> {
  // CHECK: "tfl.concatenation"(%arg0, %arg1) <{axis = 0 : i32, fused_activation_function = "NONE"}>
  %0 = "tfl.concatenation"(%arg0, %arg1) {axis = 0 : i32, fused_activation_function = "NONE"} : (tensor<1x2xi32>, tensor<1x2xi32>) -> tensor<2x2xi32>
  func.return %0 : tensor<2x2xi32>
}

// -----

func.func @testDontConvertPackToReshapeDynamicShape(%arg0: tensor<2x?xf32>) -> tensor<1x2x?xf32> {
  %0 = "tf.Pack"(%arg0) {axis = 0 : i64} : (tensor<2x?xf32>) -> tensor<1x2x?xf32>
  func.return %0 : tensor<1x2x?xf32>
  // CHECK: %[[PACK:.*]] = "tf.Pack"(%arg0) <{axis = 0 : i64}> : (tensor<2x?xf32>) -> tensor<1x2x?xf32>
  // CHECK: return %[[PACK]] : tensor<1x2x?xf32>
}

// CHECK-LABEL: while_with_id_passthrough

  %3 = stablehlo.uniform_quantize %2 : (tensor<1x2x!quant.uniform<i32:f32:1, {2.000000e+00, 2.000000e+00}>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
  return %3 : tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
}

    func.return %cst : tensor<1xi8>
  }

  // CHECK-LABEL: infer_var_handle_op_from_read
  func.func @infer_var_handle_op_from_read() -> tensor<1xi8> {
    %cst = arith.constant dense<1> : tensor<1xi8>
    %0 = "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor<!tf_type.resource<tensor<*xi8>>>
    // CHECK: "tf.VarHandleOp"() <{container = "", shared_name = "bar"}> : () -> tensor<!tf_type.resource<tensor<1xi8>>>

    // CHECK:             %[[A_SHARD:.+]] = "tf.XlaSpmdFullToShardShape"(%[[A]]) <{dim = -1 : i64, manual_sharding = "\08\03\1A\02\02\01\22\02\00\01", unspecified_dims = []}> : (tensor<2x2xi64>) -> tensor<1x2xi64>
    // CHECK:             %[[B:.+]] = "tf._XlaHostComputeMlir"(%[[A_SHARD]])
    // CHECK-SAME:          manual_sharding = true
    // CHECK-SAME:          recv_key = "host_compute_channel_0_retvals"