if (!exported_model.ok()) return exported_model.status();

        // Remove the `tpu` tag from the debug quantized saved model as it is
        // for CPU. Note the 'tpu' value should be the same as `TPU` defined in
        // tensorflow/python/saved_model/tag_constants.py.
        if (quantization_options.has_debugger_config()) {
          tags.erase("tpu");
        }
        py_function_library.SaveExportedModel(

# bazel test invocation as normal.
test:nonpip_filters --test_tag_filters=-no_oss,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310
test:nonpip_filters --build_tag_filters=-no_oss,-oss_serial,-gpu,-tpu,-benchmark-test,-v1only,-no_aarch64,-no_oss_py38,-no_oss_py39,-no_oss_py310

    mlir::ArrayAttr device_assignment_attr);

// Finds the TPU compilation device and execution devices from `devices` for a
// TPU computation subgraph. Compilation device is determined from looking up
// all TPU_SYSTEM:0 devices and choosing the CPU device associated to the first
// TPU_SYSTEM device sorted lexicographically by replica and task. Execution
// devices are determined by looking up all TPU devices associated with each

}

// "x's type V is not a named type and T is a type parameter, and x is assignable to each specific type in T's type set."
func _[
	TP0 any,
	TP1 ~_Chan,
	TP2 ~chan int | ~chan byte,
]() {
	var (
		_ TP0 = c // ERRORx `cannot use c .* as TP0 value`
		_ TP0 = C // ERRORx `cannot use C .* as TP0 value`
		_ TP1 = c
		_ TP1 = C // ERRORx `cannot use C .* as TP1 value`

#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/lib/core/status_test_util.h"

namespace mlir::quant::stablehlo {
namespace {

using ::testing::Test;

    // expected-warning @+1 {{TF2XLA TPU bridge input check: cluster op = tf.opA with cluster = cluster has successor as non cluster op tf.opB}}
    %out2, %c2 = tf_executor.island wraps "tf.opB"(%out) : (tensor<i32>) -> tensor<i32>
    // expected-error @+1 {{tf.TPUReplicatedOutput' op TF2XLA TPU bridge input check: non-cluster op = tf.opB has invalid successor op = tf.TPUReplicatedOutput}}

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {
namespace {

constexpr char kEntryFuncName[] = "main";

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {

// RUN: tf-tfrt-opt -tfrt-lower-cluster-to-runtime-ops-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:TPU:0"]} {

  // CHECK-LABEL: @converts_cluster
  func.func @converts_cluster() {
    // CHECK: %0:2 = "tf_device.launch"() <{{.*}}> ({

    arg->set_dtype(dtype);
    // TODO(lyandy): Support other arg kinds.
    if (dtype == tensorflow::DT_RESOURCE)
      arg->set_kind(tensorflow::tpu::TPUCompileMetadataProto::Arg::VARIABLE);
    else
      arg->set_kind(tensorflow::tpu::TPUCompileMetadataProto::Arg::PARAMETER);

    // Populate argument shapes.
    *arg->mutable_shape() = tensorflow::TensorShapeProto();
    if (auto ranked_tensor_type =

// TPU-annotated operations and intended to preserve backward compatibility with
// TFv1.
struct TpuV1BridgeExecutorIslandCoarsening
    : public impl::TpuV1BridgeExecutorIslandCoarseningPassBase<
          TpuV1BridgeExecutorIslandCoarsening> {
  void runOnOperation() override;
};

// Returns name of TPU cluster, if op belongs to a TPU cluster. Otherwise,
// returns `std::nullopt`.