auto [partitioned_output, id_ops] =
            GetSingleUserOfType<mlir::TF::TPUPartitionedOutputV2Op>(
                per_replica_result);
        if (partitioned_output) {
          erase_list.insert(erase_list.end(), id_ops.begin(), id_ops.end());
          partitioned_outputs[cluster_result_id].emplace_back(
              partitioned_output);
        }
      }
    }
  }

// TPUPartitionedOutputV2 ops. This test ensures that they are erased
// and not considered within the clustered computation. It also ensures that
// the expected interleaving pattern is present in the output.

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>)

      status = ReplacePartitionedOp(num_cores_per_replica, partitioned_input);
    } else if (auto partitioned_output =
                   llvm::dyn_cast_or_null<TF::TPUPartitionedOutputOp>(op)) {
      status = ReplacePartitionedOp(num_cores_per_replica, partitioned_output);
    }

    if (status.has_value()) {
      if (failed(*status) || !op->use_empty()) return WalkResult::interrupt();

    mlir::TF::TPUPartitionedOutputV2Op partitioned_output;
    for (auto user : old_parallel_execute_output.getUsers()) {
      if (auto partitioned_output_user =
              llvm::dyn_cast_or_null<mlir::TF::TPUPartitionedOutputV2Op>(
                  user)) {
        partitioned_output = partitioned_output_user;
        break;
      }
    }
    if (partitioned_output) {
      if (!old_parallel_execute_output.hasOneUse())

}

// Returns the op sharding attribute from a partitioned operator.
std::optional<StringRef> GetXlaShardingFromOperator(mlir::Operation* op) {
  if (auto partitioned_output =
          llvm::dyn_cast<mlir::TF::TPUPartitionedOutputV2Op>(op)) {
    return partitioned_output.get_XlaSharding();
  } else if (auto partitioned_input =
                 llvm::dyn_cast<mlir::TF::TPUPartitionedInputV2Op>(op)) {

    %partitioned_output:2 = "tf.TPUPartitionedOutputV2"(%computation) {N = 2 : i64, partition_dims = []} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
    "tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
    "tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()

    }

    auto partitioned_output = builder.create<TF::TPUPartitionedOutputV2Op>(
        cluster_func->getLoc(), partitioned_output_types, result,
        partitioned_input.getPartitionDimsAttr(),
        partitioned_input.get_XlaShardingAttr());
    for (auto [i, value] : llvm::enumerate(partitioned_output.getOutput())) {
      const auto& resource = packed_input ? inputs[0] : inputs[i];

  %partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {N = 2 : i64, _XlaSharding = "", partition_dim = -1 : i64} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
  "tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
  "tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()

      %partitioned_output:2 = "tf.TPUPartitionedOutput"(%computation) {N = 2 : i64, _XlaSharding = "", partition_dim = -1 : i64} : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
      "tf.AssignVariableOp"(%arg0, %partitioned_output#0) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()
      "tf.AssignVariableOp"(%arg1, %partitioned_output#1) : (tensor<!tf_type.resource<tensor<i32>>>, tensor<i32>) -> ()