// exportable and verify some properties of the Tensorflow Executor MLIR that
// are required by Graph Export. It will return an error if it cannot.
//
// Input: A MLIR Module in the Tensorflow Dialect with no
// `tf_device.cluster_func` ops.
// Output: A MLIR module in the Tensorflow Executor Dialect.
tensorflow::Status ExportFromTensorflowDialectToExecutor(
    mlir::ModuleOp module, llvm::StringRef module_name = llvm::StringRef());

}  // namespace v2

            tf_device.return %4#0, %4#1 : tensor<2048xi32>, tensor<2048xi32>
            }) {device = "TPU_REPLICATED_HOST_0"} : () -> (tensor<2048xi32>, tensor<2048xi32>)
      %1 = "tf_device.cluster_func"(%0#0, %0#1) {_replication_info = "cluster_test_fn", func = @tpu_func} : (tensor<2048xi32>, tensor<2048xi32>) -> tensor<2048xi32>
      return %1: tensor<2048xi32>
   }
   // CHECK-LABEL: func @tpu_func

// exportable and verify some properties of the Tensorflow Executor MLIR that
// are required by Graph Export. It will return an error if it cannot.
//
// Input: A MLIR Module in the Tensorflow Dialect with no
// `tf_device.cluster_func` ops.
// Output: A MLIR module in the Tensorflow Executor Dialect.

ABSL_DEPRECATED(
    "Use v2/tf_dialect_to_executor.h::ExportFromTensorflowDialectToExecutor "
    "instead.")

}

void TPUAnnotateDynamicShapeInputsPass::runOnOperation() {
  getOperation().walk([&](tf_device::ClusterFuncOp cluster_func_op) {
    Builder builder(cluster_func_op->getContext());
    // Skip non-tpu device cluster_func.
    auto cluster_id =
        cluster_func_op->getAttrOfType<StringAttr>(TF::kReplicationInfoAttr);
    if (!cluster_id) return WalkResult::advance();

    llvm::SmallVector<int, 4> dynamic_shape_arg_index;

  ReplaceClusterReturnWithReturn(return_op, builder);

  symbol_table->insert(outlined_func);
  return outlined_func;
}

// Outlines body of `tf_device.cluster` into a function and create a
// `tf_device.cluster_func` to invoke that function. `tf_device.cluster` is
// removed afterwards.`
void OutlineCluster(tf_device::ClusterOp cluster_op, SymbolTable* symbol_table,
                    OpBuilder* builder) {

}

// -----

func.func @parallel_execute_regions_with_data_results(%arg0: tensor<i32>) -> tensor<i32> {
  // expected-error @+1 {{'func' attribute refers to an undefined function: undefined_func}}
  %0 = "tf_device.cluster_func"(%arg0) {func = @undefined_func} : (tensor<i32>) -> tensor<i32>
  func.return %0 : tensor<i32>

  let summary = "Invokes a program via IFRT on a device";

  let description = [{
    This op calls an IFRT program uniquely identified by the given program id.

    During lowering from a `tf_device.cluster_func` op to a `tf.IfrtCall` op,
    the region owned by the former will be outlined to a function with a
    `tfrt_ifrt_serving.program_id` attribute. After that, the runtime ensures

    %1 = "tf.Shape"(%arg1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = ""} : (tensor<*x!tf_type.string>) -> tensor<*xi32>
    %2:3 = "tf.UnpackHyp"(%arg1) {_tpu_replicate = "cluster_Fn", _xla_outside_compilation = "0", device = "", max_seq_length = 16 : i64} : (tensor<*x!tf_type.string>) -> (tensor<*xi32>, tensor<*xi32>, tensor<*xf32>)