location, cluster_func_output_index, output_sharding_config,
      cluster_to_core_index, cluster_idx, new_parallel_execute,
      &outputs_to_merge);
  if (failed(status)) return mlir::failure();

  // Creates a tree of Concat ops that merges outputs from multiple logical
  // devices to a single replica output.
  const xla::OpSharding& sharding =
      output_sharding_config[cluster_func_output_index];

// different outputs depending on the output sharding configuration.
mlir::LogicalResult GetOutputTypesForLogicalDeviceComputation(
    int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config,
    mlir::tf_device::ClusterFuncOp cluster_func,
    llvm::SmallVectorImpl<mlir::Type>* output_types,
    llvm::SmallVectorImpl<int>* cluster_to_core_index);

}

// Creates a `tf.TPUExecute` op that executes TPU program.
LogicalResult BuildExecuteOp(
    const int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config,
    llvm::ArrayRef<Value> inputs, tf_device::ClusterFuncOp cluster_func,
    OpBuilder* builder, TF::TPUExecuteOp* execute_op) {
  // TODO(b/139377366): Need to snapshot all resource variable inputs in