func.func @add_with_activation_transpose_rank_two(%arg0: tensor<1x2xf32>) -> tensor<2x1xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<2x1xf32>
  %1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
  %2 = stablehlo.add %1, %0 : tensor<2x1xf32>
  return %2 : tensor<2x1xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose

  %0 = "mhlo.reshape"(%arg0) : (tensor<1x1x512xf32>) -> tensor<1x512xf32>
  %1 = "mhlo.dot"(%0, %arg1) : (tensor<1x512xf32>, tensor<512x13x!quant.uniform<i8:f32, 0.00285>>) -> tensor<1x13xf32>
  %2 = "mhlo.reshape"(%1) : (tensor<1x13xf32>) -> tensor<1x1x13xf32>
  func.return %2 : tensor<1x1x13xf32>

// CHECK:      %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME:   dot_dimension_numbers = #mhlo.dot<

  %2 = "tfl.relu"(%arg0) : (tensor<1xf32>) -> tensor<1xf32>
  // CHECK: tac.device = "CPU", tac.inference_type = "FLOAT"
  %3 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return
}

func.func @notAnnotateConst(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
  // CHECK-NOT: tac.device tac.inference_type

  %split_dim = arith.constant dense<0> : tensor<i32>
  // expected-error @+1 {{'tfl.split' op output #0 should be 'tensor<8x4xf32>'}}
  %0, %1 = "tfl.split"(%split_dim, %arg0) {num_splits = 2 : i32} : (tensor<i32>, tensor<16x4xf32>) -> (tensor<8x2xf32>, tensor<8x2xf32>)
  func.return %0, %1 : tensor<8x2xf32>, tensor<8x2xf32>
}

// -----

For example, if we have the code

```mlir
  %0 = "tf.Const"() {value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
  %1 = "tf.Const"() {device = "", value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
  %2 = "tf.Const"() {device = "baz", value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
```

then running this pass with 'default-device=foobar', we get:

```mlir

  %2 = "tf.Transpose"(%1, %cst_0): (tensor<1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  func.return %2 : tensor<2x1xf32>

// CHECK: %cst = arith.constant
// CHECK: %[[trans:.*]] = "tf.Transpose"
// CHECK-SAME: -> tensor<2x1xf32>
// CHECK: %[[q:.*]] = "tfl.quantize"(%[[trans]]) <{qtype = tensor<2x1x!quant.uniform<u8:f32, 1.000000e+00>>}>

    }) {is_stateless = false} : (tensor<i32>, tensor<!tf_type.variant<tensor<?x1xf32>>>) -> (tensor<i32>, tensor<!tf_type.variant<tensor<?x1xf32>>>)
    %elem_1 = "tf._SomeOtherOp"() : () -> tensor<8x1xf32>
    %tl_set_item = "tf.TensorListSetItem"(%while#1, %one, %elem_1) : (tensor<!tf_type.variant<tensor<?x1xf32>>>, tensor<i32>, tensor<8x1xf32>) -> tensor<!tf_type.variant<tensor<?x1xf32>>>
    func.return
  }

func.func @torch_index_select(%arg0: tensor<2x1xf32>, %arg1: tensor<2xi32>) -> tensor<2x1xf32> {
  %0 = "mhlo.torch_index_select"(%arg0, %arg1) {
    batch_dims = 0 : i64, dim = 0 : i64
  } : (tensor<2x1xf32>, tensor<2xi32>) -> tensor<2x1xf32>
  func.return %0 : tensor<2x1xf32>
}

// CHECK-LABEL:   func @lowered_cumsum(
// CHECK-SAME:      %[[VAL_0:.*]]: tensor<4x12xf32>) -> tensor<4x12xf32> {

mlir_module = '''python
func @main(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10x10xf32> {
   %add = "magic.op"(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10x10xf32>
   return %ret : tensor<10x10xf32>
}
'''

@tf.function
def foo(x, y):
  return mlir_passthrough_op([x, y], mlir_module, Toutputs=[tf.float32])