operands.append(read_operands.begin(), read_operands.end());

    auto loc = cluster_func.getLoc();
    auto new_cluster_func = builder.create<tf_device::ClusterFuncOp>(
        loc, cluster_func.getResultTypes(), operands, cluster_func->getAttrs());
    cluster_func.replaceAllUsesWith(new_cluster_func);
    func::FuncOp func = cluster_func.getFuncOp();
    Block& block = func.front();

      cluster_func->getParentOfType<tf_device::ParallelExecuteOp>();
  if (old_parallel_execute &&
      cluster_func->getParentOp() != old_parallel_execute) {
    cluster_func->emitError() << "The ParallelExecute ancestor of a "
                                 "ClusterFunc must be its direct parent.";
    return failure();
  }
  if (!old_parallel_execute)
    old_parallel_execute = TF::BuildParallelExecuteOp(cluster_func, builder);

  // CHECK:      tf_device.cluster_func
  // CHECK-SAME: input_sharding_configuration = ["\01\02\03"]
  // CHECK-SAME: output_sharding_configuration = []
  "tf_device.cluster_func"(%1) {func = @cluster_func, use_spmd_for_xla_partitioning = true, num_cores_per_replica = 1 : i64} : (tensor<*xf32>) -> ()
  func.return
}

// CHECK-LABEL: func @cluster_func

void MarkInputOutputAliasesPass::runOnOperation() {
  SmallVector<tf_device::ClusterFuncOp, 4> cluster_funcs;
  ModuleOp module = getOperation();
  module.walk([&](tf_device::ClusterFuncOp cluster_func) {
    // Map resource values to pair of input-output indices.
    llvm::DenseMap<Value, AliasInfo> resource_alias_info_map;
    if (failed(BuildAliasingInfo(cluster_func, resource_alias_info_map)) ||
        resource_alias_info_map.empty()) {
      return;

}

// Extracts input/output sharding configuration of `cluster_func` by parsing
// XlaSharding ops inside the `cluster_func`.
LogicalResult IdentifyXlaShardingForTPUComputation(
    Builder* builder, mlir::tf_device::ClusterFuncOp cluster_func) {
  // Look up function definition from module.
  mlir::func::FuncOp func =
      cluster_func->getParentOfType<ModuleOp>()

    %1 = "tf.A"() : () -> tensor<?xi32>
    tf_device.return %1 : tensor<?xi32>
  }) {_cluster_outlined_function_name = "cluster_func"} : () -> tensor<?xi32>
  func.return %0 : tensor<?xi32>
}

// CHECK: "tf_device.cluster_func"

    mlir::tf_device::ClusterFuncOp cluster_func,
    mlir::SmallVector<xla::OpSharding, 4>* output_sharding_list) {
  output_sharding_list->reserve(cluster_func.getNumResults());

  const auto output_sharding_attrs =
      cluster_func.getOperation()->getAttrOfType<mlir::ArrayAttr>(
          kOutputShardingAttr);
  if (!output_sharding_attrs)
    return cluster_func.emitError(

  // CHECK-NEXT: [[READ:%.*]] = "tf.ReadVariableOp"([[ARG2]])
  // CHECK-NEXT: [[CLUSTER:%.*]]:2 = "tf_device.cluster_func"([[ARG0]], [[ARG1]], [[READ]])
  // CHECK-SAME: _replication_info = "write", _xla_compile_device_type = "TPU"
  %0:2 = "tf_device.cluster_func"(%arg0, %arg1) {_replication_info = "write", _xla_compile_device_type = "TPU", func = @write_func} : (tensor<i32>, tensor<f32>) -> (tensor<f32>, tensor<i32>)

  %1 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
  // CHECK:      [[COMPUTATION:%.+]] = "tf_device.cluster_func"([[INPUT]])
  %2 = "tf_device.cluster_func"(%1) {func = @computation, use_spmd_for_xla_partitioning = true} : (tensor<i32>) -> tensor<i32>
  // CHECK:      [[OUTPUT:%.+]]:2 = "tf.TPUPartitionedOutputV2"([[COMPUTATION]])
  // CHECK-SAME: _XlaSharding = ""

  let summary = "Rewrites a `tf_device.cluster_func` on TPUs into TPU runtime operations.";

  let description = [{
    This pass rewrites a `tf_device.cluster_func` operation into a sequence of `tf._TPUCompileMlir`
    and `tf.TPUExecute` operations. `tf._TPUCompileMlir` contains a MLIR module that is
    functionally equivalent to the function referenced by `tf_device.cluster_func`.