// RUN: tf-quant-opt %s -quant-convert-tpu-model-to-cpu -inline -quant-cast-bf16-ops-to-f32 -split-input-file | \
// RUN: FileCheck %s

// Remove TPU related ops.
func.func @tpu_conv(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x2x2xf32> {
  %0 = "tf.TPUOrdinalSelector"() {device = ""} : () -> tensor<?xi32>
  %1 = "tf.TPUPartitionedCall"(%arg0, %0) {autotuner_thresh = 0 : i64, device = "", f = @tpu_func_0_optim0} : (tensor<1x3x4x3xf32>, tensor<?xi32>) -> tensor<1x3x2x2xf32>

  // CHECK: %[[IDENTIFY:.*]] = "tf.Identity"(%[[SUBGRAPH_0]]#1) {device = ""} : (tensor<1024x3xf32>) -> tensor<1024x3xf32>
  // CHECK: %[[SUBGRAPH_1:.*]] = "tf.XlaCallModule"() <{Sout = [#tf_type.shape<1024x3>], {{.*}} ["CPU", "TPU"], {{.*}}}> {_entry_function = @_stablehlo_main_1

using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::OpPassManager;
using mlir::PassManager;
using mlir::func::FuncOp;

// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
tensorflow::Status RunTFXLABridge(
    ModuleOp module,
    llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,
    llvm::StringRef module_name = llvm::StringRef(),

// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-non-tpu -split-input-file -verify-diagnostics %s | FileCheck %s

module attributes {tf.versions = {producer = 888 : i32}, tf.devices = ["/job:worker/replica:0/task:0/device:CPU:0", "/job:worker/replica:0/task:0/device:TPU_SYSTEM:0", "/job:worker/replica:0/task:0/device:GPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: "tf.XlaLaunch"()

                          WeightOnlyPtqComponent::kName, *function_aliases,
                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

  return builder->create<mlir::TF::ConcatOp>(
      location, output_type, concat_dimension_op.getOutput(), inputs);
}

// For tile sharded inputs to TPU computation, inject split op between the
// input values and TPU computation so that tiled input values are passed in
// as inputs to TPU computations. If more than one dimension is sharded, then
// a tree of connected split ops are added before tf_device.parallel_execute op.

                          PostCalibrationComponent::kName, *function_aliases,
                          *ctx, *module));

  // Remove the `tpu` tag for exporting because the output quantized model is
  // essentially a CPU model.
  tags.erase("tpu");

  py_function_library.SaveExportedModel(
      dst_saved_model_path, post_calibrated_exported_model,
      src_saved_model_path, tags, signature_def_map);

    "tf_device.cluster_func"(%arg0) <{func = @empty_func}> {_dynamic_arg_index = [], _replication_info = "cluster", _xla_compile_device_type = "TPU", allow_soft_placement = false, computation_shape = [], device = "", device_assignment = [0, 0, 0, 0, 0, 0, 0, 1], host_compute_core = [], input_sharding_configuration = ["{devices=[2,1]0,1}"], num_cores_per_replica = 2 : i64, output_sharding_configuration...

// RUN: tf-opt %s -tf-tpu-colocate-composite-resource-ops | FileCheck %s

// Tests ReadVariable op using composite device resource is wrapped inside
// tf_device.Cluster.

// CHECK-LABEL: func @testReadVariableOpColocated
// CHECK-SAME: (%[[ARG0:.*]]: tensor<*x!tf_type.resource<tensor<4xf32>>>)
func.func @testReadVariableOpColocated(%arg0: tensor<*x!tf_type.resource<tensor<4xf32>>>) {
  // CHECK:      tf_device.replicate

// RUN: tf-opt %s -allow-unregistered-dialect --tf-move-tpu-compile-to-front --split-input-file | FileCheck %s

module {

// CHECK-LABEL: does_basic_reordering
func.func @does_basic_reordering() -> () {
   // CHECK: _TPUCompileMlir
   // CHECK-SAME: X
   // CHECK: _TPUCompileMlir
   // CHECK-SAME: Y
   // CHECK: OpA
   // CHECK: OpB
   // CHECK: OpC
   "tf.OpA"() : () -> ()