// CHECK: %[[READ_VAR_1:[0-9]*]] = "tf.ReadVariableOp"(%arg1)
    %read1 = "tf.ReadVariableOp"(%arg1) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
    // CHECK-NOT: tf.TPUPartitionedInputV2
    %partitioned_input = "tf.TPUPartitionedInputV2"(%read0, %read1) {N = 2 : i64, partition_dims = []} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    // CHECK: %[[COMPILE_OUTPUT:[0-9]*]]:3 = "tf_device.launch"

  func.return %0 : tensor<1xi1>
}

// -----

// Test invalid tf.TPUPartitionedInputV2 with packing
func.func @testPackedTPUPartitionedInputV2(tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<4x4xf32> {
^bb0(%arg0: tensor<2x4xf32>, %arg1: tensor<2x4xf32>):
  // expected-error @+1 {{expected 1 inputs, got 2}}

      tensorflow::GetTypeFromTFTensorShape({}, IntegerType::get(context, 1)));
  return success();
}

//===----------------------------------------------------------------------===//
// TPUPartitionedInputV2
//===----------------------------------------------------------------------===//

// This method mimics this op's core/TF-level shape inference logic
LogicalResult TPUPartitionedInputV2Op::verify() {