"//tensorflow/compiler/tf2xla/kernels:xla_ops",
        "//tensorflow/core:framework",
        "//tensorflow/core:test_main",
        "//tensorflow/core/protobuf/tpu:compile_metadata_proto_cc",
        "//tensorflow/core/tpu/kernels/xla:host_compute_ops",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings:string_view",
        "@llvm-project//mlir:FuncDialect",
        "@llvm-project//mlir:IR",

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-rewrite-ifrt-load-variable %s | FileCheck %s

// Variable is used by both CPU and TPU
//
// CHECK-LABEL: func @serving_default(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    [[HANDLE:%.*]] = "tf.VarHandleOp"()
// CHECK-NEXT:    [[ARRAYKEY:%.*]], [[FURTURE:%.*]] = "tf_mlrt.tf_ifrt_load_variable"([[HANDLE]])

    return %0 : tensor<i1>
  }
  // Generated functions
  // non_tpu should have to TPU ops - just identity and return (in this test).
  // CHECK: func.func private @_func_non_tpu
  // CHECK-NEXT: tf.Identity
  // CHECK-NEXT: return

  // sc_forward should have TPU ops including replicated outputs but not inputs
  // CHECK: func.func private @_func_sc_forward
  // CHECK-NOT: TPUReplicatedInput

    %cst_12 = "tf.Const"() {_replication_info = "repl_info", _xla_compile_device_type = "TPU", device = "", value = dense<1> : tensor<i64>} : () -> tensor<i64>
    "tf.AssignAddVariableOp"(%37, %cst_12) {_has_manual_control_dependencies = true, _replication_info = "while/cluster_while_body_451", _xla_compile_device_type = "TPU", device = ""} : (tensor<*x!tf_type.resource<tensor<i64>>>, tensor<i64>) -> ()

#include "llvm/ADT/DenseSet.h"
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Support/TypeID.h"  // from @llvm-project

namespace mlir {
namespace TF {

// A global ops registry that is used to hold TPU embedding ops.
//
// Example:
//    TPUEmbeddingOpsRegistry::Global().Add<TF::FooOp>();
//    for (auto op_type_id : TPUEmbeddingOpsRegistry::Global().GetOpsTypeIds())
//    {
//      ...
//    }

  // For all other inputs that go from TPU op to TPU op, insert the
  // TPUOutput/Input pair.

  // Add TPUReplicatedInput/TPUReplicatedOutput pairs along each edge.
  // TODO(bfontain): Should be merged with the above loop.
  llvm::SetVector<Value> values_to_add_nodes;

  for (Operation* op : backward_pass_ops) {
    // TODO(bfontain): Should validate that all the TPU ops are in the same
    // replication region.

// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-tpu-colocate-splits | FileCheck %s

// CHECK-LABEL: func @colocate_split_with_pred
func.func @colocate_split_with_pred() {
  // CHECK: Split
  // CHECK-SAME: _class = ["loc:@class"]
  tf_executor.graph {
    %c, %control0 = tf_executor.island wraps "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>

// RUN: tf-opt %s -tf-executor-tpu-v1-island-inlining | FileCheck %s

// Check that the nested module is inlined and erased.

module {
// CHECK-LABEL: func @func0
  func.func @func0(%arg0: tensor<i1>) -> tensor<f32> {
    %0 = tf_executor.graph {
// CHECK-NOT: PartitionedCall
// CHECK: "tf.opA"

    const FunctionLibraryRuntime& function_library, DeviceBase* device,
    se::Stream* stream, const XlaPlatformInfo& platform_info,
    bool has_ref_vars);

// Returns created options for XLA compiler when TFRT-TPU is used.
XlaCompiler::Options GenerateCompilerOptionsForTfrtTpu(
    const DeviceCompiler<xla::LocalExecutable, xla::LocalClient>&
        xla_device_compiler,
    const FunctionLibraryRuntime& function_library);

//
// Duplicating shape-determining constants is required to place constants that
// affect the shape of a tensor to be placed in the TPU graph instead of in the
// CPU graph, when the graph gets converted for TPU inference. This allows these
// constants to be known at XLA compilation time.
void AddExportPasses(mlir::PassManager& pm,
                     bool duplicate_shape_determining_constants);