// CHECK:      [[INPUT0:%.+]] = "tf.TPUPartitionedInputV2"([[READ0]], [[READ1]])
  // CHECK-DAG:  [[READ2:%.+]] = "tf.ReadVariableOp"([[ARG0]])
  // CHECK-DAG:  [[READ3:%.+]] = "tf.ReadVariableOp"([[ARG1]])
  // CHECK:      [[INPUT1:%.+]] = "tf.TPUPartitionedInputV2"([[READ2]], [[READ3]])

  // CHECK: [[PI_0:%.*]] = "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])
  %pi_0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {device = "", partition_dims = []} : (tensor<!tf_type.resource<tensor<10x3xf32>>>, tensor<!tf_type.resource<tensor<10x3xf32>>>) -> tensor<!tf_type.resource<tensor<10x3xf32>>>
  // CHECK: [[PI_1:%.*]] = "tf.TPUPartitionedInputV2"([[ARG2]], [[ARG3]])

  // CHECK: [[PI_0:%.*]] = "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])
  // CHECK-SAME: _XlaSharding = ""
  // CHECK-SAME: partition_dims = []
  // CHECK: [[PI_1:%.*]] = "tf.TPUPartitionedInputV2"([[ARG2]], [[ARG3]])
  // CHECK-SAME: _XlaSharding = ""
  // CHECK-SAME: partition_dims = []

        return llvm::isa_and_nonnull<TF::TPUPartitionedInputV2Op>(
            input.getDefiningOp());
      }))
    return replicated_input.emitOpError()
           << "expects all inputs from 'tf.TPUPartitionedInputV2' ops";

  const auto metadata_iter =
      replicated_input->getBlock()->getOps<TF::TPUReplicateMetadataOp>();
  TF::TPUReplicateMetadataOp metadata;
  if (!metadata_iter.empty()) metadata = *(metadata_iter.begin());

      // ClusterFunc's use of the Read is replaced with use of the
      // TPUPartitionedInputV2.
      read_use.set(new_partitioned_input);
    } else {
      // Outside compiled code's use of the Read after TPUPartitionedInputV2 is
      // replaced with use of the first Read before the TPUPartitionedInputV2.
      if (sharding.type() != xla::OpSharding::REPLICATED) {

func.func @partitioned_outputs(%arg0: tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>, tensor<?xi32>, tensor<?xi32>) {
  %pi0 = "tf.TPUPartitionedInputV2"(%arg0) {N = 2, partition_dims = [], _XlaSharding = "", is_packed = true} : (tensor<?xi32>) -> (tensor<?xi32>)
  %pi1 = "tf.TPUPartitionedInputV2"(%arg0) {N = 2, partition_dims = [], _XlaSharding = "", is_packed = true} : (tensor<?xi32>) -> (tensor<?xi32>)

// replicate sharding ("").

// CHECK-LABEL: func @partitioned_input_output
func.func @partitioned_input_output(%arg0: tensor<*xi32>, %arg1: tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>) {
  %0 = "tf.TPUPartitionedInputV2"(%arg0) {_XlaSharding = "\01\02\03", partition_dims = []} : (tensor<*xi32>) -> tensor<*xi32>
  // CHECK:      tf_device.cluster_func
  // CHECK-SAME: input_sharding_configuration = ["\01\02\03", ""]

          }
        }
      }
    }
  }
  for (auto cluster_operand : cluster.getOperands()) {
    Operation* def = cluster_operand.getDefiningOp();
    // This pass assumes that a TPUPartitionedInputV2 is preceeded by
    // ReadVariable ops, and not vice versa. An earlier pass,
    // TPUResourceReadsWritesPartitioning, should have ensured this
    // precondition.
    if (!def) continue;

          }
        }
        // When model parallelism is used in conjunction with data parallelism
        // for resource inputs, we need to collect the per replica resource
        // inputs from input to `tf.TPUPartitionedInputV2` ops.
        if (auto pi = llvm::dyn_cast_or_null<mlir::TF::TPUPartitionedInputV2Op>(
                def)) {
          if (pi->getNumOperands() != num_cores_per_replica)
            status = pi.emitOpError()

    // 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"