getOperation()->getParentOfType<ModuleOp>(), &devices)))
    return signalPassFailure();
  for (const auto& device_name : devices.device_names()) {
    if (device_name.has_type && device_name.type == "TPU") ++device_num;
  }

  if (device_num == 0) return;

  llvm::SmallVector<Operation*, 4> require_device_ordinal_ops;
  getOperation().walk([&](Operation* op) {
    if (RequiresDeviceOrdinalAttribute(op)) {

  void runOnFunction(
      func::FuncOp func,
      const TF::ResourceAliasAnalysis::Info& resource_alias_analysis);

  StringRef getArgument() const final { return "tf-tpu-dynamic-layout-pass"; }

  StringRef getDescription() const final {
    return "Inserts TPU layout ops to determine layout at run time.";
  }
};

// Checks if the input producer op is supported in this transform. Right now, we

    return absl::OkStatus();
  }

  // TFRT-TPU is used if device type is `DEVICE_TPU` and platform_info does not
  // have `xla_device_metadata`. This is used for TFRT-TPU when
  // BuildXlaDeviceCompiler() is called in GetCompilerIr(). Currently only
  // lowering to HLO is needed there and xla::LocalClient doesn't support
  // building the executable for TFRT-TPU and hence, is set to nullptr here.

  func.func @main() attributes {tf.entry_function = {inputs = "", outputs = ""}} {
    tf_executor.graph {
      %0:2 = tf_executor.island wraps "tf.Const"() {device = "TPU:0", name = "const", dtype = "tfdtype$DT_INT32", value = dense<[1, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
      tf_executor.fetch
    }
    func.return
  }
}

// CHECK-NOT: node

// CHECK: library

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

  let summary = "Op that compiles a computation in MLIR into a TPU program, and loads and executes it on a TPU device.";

  let description = [{
For the internal use of the TPU compiler.

'static_shapes' are tensors specifying the maximum dimension sizes for the tensors specified in `dynamic_operands`.
'args' are inputs to the TPU computation.

  // 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"()

*   Compiler researchers and implementers looking to optimize performance and
    memory consumption of machine learning models
*   Hardware makers looking for a way to connect their hardware to TensorFlow,
    such as TPUs, portable neural hardware in phones, and other custom ASICs
*   People writing language bindings that want to take advantage of optimizing
    compilers and hardware acceleration.