%1 = "tf.StatefulPartitionedCall"(%arg0) {config = "", config_proto = "", executor_type = "", f = @stateful_partitioned_call_func} : (tensor<20xi32>) -> tensor<*xi32>
    // CHECK: tf.TPUPartitionedCall
    // CHECK-SAME: (tensor<20xi32>, tensor<?xi32>) -> tensor<20xi32>
    %2 = "tf.TPUPartitionedCall"(%arg0, %arg1) {autotuner_thresh = 0 : i64, f = @tpu_partitioned_call_func} : (tensor<20xi32>, tensor<?xi32>) -> tensor<*xi32>

  func.return %0#0 : tensor<2xi32>
}

// -----

func.func @unpack(%arg0: tensor<2x3xi32>) -> tensor<2xi32> {
  // expected-error @+2 {{failed to infer returned types}}
  // expected-error @+1 {{op inferred type(s) 'tensor<2xi32>', 'tensor<2xi32>', 'tensor<2xi32>' are incompatible with return type(s) of operation 'tensor<2xi32>', 'tensor<2x1xi32>', 'tensor<2xi32>'}}

  %s0 = "tf.Const"() {value = dense<[4, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
  %s1 = "tf.Const"() {value = dense<[4, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
  // expected-error @+1 {{requires dimension 0 size of 'r0' to be 1 but got 0}}
  %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) : (tensor<2xi32>, tensor<2xi32>) -> (tensor<0xi32>, tensor<0xi32>)
  func.return %r0, %r1 : tensor<0xi32>, tensor<0xi32>
}

  func.return %2, %3, %4, %5: tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>
}

// CHECK-LABEL: func @addNWithZerosInt
func.func @addNWithZerosInt(%arg0: tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) {
  %0 = "tf.Const"() {value = dense<1> : tensor<2xi32>} : () -> tensor<2xi32>
  %1 = "tf.Const"() {value = dense<0> : tensor<2xi32>} : () -> tensor<2xi32>

// CHECK-DAG:       [[VAL_33:%.*]] = arith.constant dense<[-1, 0, 0]> : tensor<3xi32>
// CHECK-DAG:       [[VAL_34:%.*]] = arith.constant dense<0> : tensor<3xi32>

  %0 = "tf.Pack"(%arg0, %arg1, %arg2) {axis = 1 : i64} : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x3xi32>
  func.return %0 : tensor<2x3xi32>

// CHECK-LABEL: pack3Tensors
// CHECK: "tfl.pack"(%arg0, %arg1, %arg2) <{axis = 1 : i32, values_count = 3 : i32}> : (tensor<2xi32>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x3xi32>
}

        %5:2 = "tf.C"(%4) : (tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi32>)
        tf_device.return %5#0, %5#1 : tensor<2xi32>, tensor<2xi32>
      }) {num_cores_per_replica = 1, topology =  "", device_assignment =  []} : () -> (tensor<2xi32>, tensor<2xi32>)
      tf_device.return %2#1 : tensor<2xi32>
    }
    func.return %1 : tensor<2xi32>
  }

// CHECK-DAG: %[[START:.+]] = arith.constant dense<{{\[1, 2\]}}> : tensor<2xi32>
// CHECK-DAG: %[[SIZE:.+]] = arith.constant dense<2> : tensor<2xi32>
// CHECK: %[[SLICE:.+]] = "tfl.slice"(%[[ARG0]], %[[START]], %[[SIZE]]) : (tensor<3x4x!quant.uniform<i8:f32, 2.000000e+00:-1>>, tensor<2xi32>, tensor<2xi32>) -> tensor<2x2x!quant.uniform<i8:f32, 2.000000e+00:-1>>
// CHECK: return %[[SLICE]]

// -----

    ```mlir
      %2 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
      %3 = "tf.B"(%2) {device = "tpu0"} : (tensor<?xi32>) -> tensor<?xi32>
      %4 = "tf.C"(%2, %3) {device = "tpu0"} : (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
      %5 = "tf.D"(%4) : (tensor<?xi32>) -> tensor<?xi32>
    ```

    After the pass, we will have:

    ```mlir
      %0 = "tf.A"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>

  // The first Operand is assumed to be a TensorType around a variant with a
  // single subtype (e.g. tensor<!tf_type.variant<tensor<2xi32>>>). We will
  // copy this type to the first result, and copy the singular variant subtype
  // to the second result (tensor<2xi32>).
  DCOMMENT_OP(op, "Inferring shape for TensorListPopBackOp.");

  auto src_list_handle_t =