%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>) -> ()

// RUN: tf-opt %s -tf-tpu-device-propagation | FileCheck %s

// Tests function passthrough values.

// CHECK-LABEL: func @testArgToRet
// CHECK-SAME: ({{%.+}}: tensor<i64> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"})
// CHECK-SAME: -> (tensor<i64> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"})
func.func @testArgToRet(%arg0: tensor<i64> {tf.device = "/job:localhost/replica:0/task:0/device:TPU:0"}) -> tensor<i64> {

  // of dropping the attribute. This need not be correct for models not trained
  // on TPU.
  // Extracts the StableHLO module from tf.XlaCallModuleOp if the StableHLO
  // module is serialized in it.
  pm.addPass(mlir::odml::CreateLegalizeTFXlaCallModuleToStablehloPass());

  // Preprocesses TPU-targeting StableHLO module for support in TF Quantizer.
  pm.addPass(mlir::quant::CreateConvertTpuModelToCpuPass());

        // device. Device attribute is used to infer the appropriate sharding
        // within TPUs for this op.
        // TODO(b/183598857): Use explicit sharding ops from the front-end.
        // For example, dequeue ops generated by
        // tensorflow/python/tpu/tpu_feed.py
        if (!tensorflow::IsTPUReplicatedCore(attr.getValue())) {
          inner_op->removeAttr(kDeviceAttr);
        }

  return str_attr &&
         (str_attr.getValue().find("COMPOSITE") != llvm::StringRef::npos);
}
}  // namespace

// Rewrites unpartitioned resource reads and writes to partitioned resource
// reads and writes. The TPU computation from the frontend is generated in such
// a way that resource operations operate on the unpartitioned resource handle
// (from a `tf.TPUReplicatedInput`). This results in resource reads and writes

// and logical core (`logical_core`).
// `replica_id` is the index of the ancestor ReplicateOp in [0, num_replicas).
// `logical_core` is the index of the TPU core in [0, num_cores_per_replica).
// `device_ordinal` is the index of the TPU core relative to its host.
LogicalResult GetDeviceOrdinal(const std::optional<DictionaryAttr>& devices,
                               const unsigned replica_id,

          func.getLoc(), op.getOperandTypes(), /*dynamic_key=*/dynamic_key,
          op.getSendKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),
          rewriter.getStringAttr("TPU"));
      for (auto result :
           llvm::zip(cloned_func.getArguments(), recv_at_host->getResults())) {
        std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
      }

// RUN: tf-tfrt-opt --split-input-file -pass-pipeline='builtin.module(pre-parallel-tf-to-mlrt{use-tpu-host-allocator-for-inputs=true},tf-mlrt-parallelization{tfrt-cost-threshold=4},tf-to-mlrt)'  %s | FileCheck %s --dump-input=fail --dump-input-filter=all

func.func @callee(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>) {
  func.return %arg0: tensor<i32>
}

// CHECK-LABEL: func @batch_function
func.func @batch_function(%arg0: tensor<i32>) -> (tensor<i32>) {

  auto status_or_device_coodinates =
      tensorflow::GetDeviceCoordinates(device_assignment_attr);
  if (!status_or_device_coodinates.ok())
    return cluster.emitError()
           << "error in fetching tpu device coordinates: "
           << status_or_device_coodinates.status().message();

  status_or_tpu_device_assignment =
      tensorflow::GetTPUCompilationAndExecutionDevices(

                                 RewritePatternSet* patterns) {
  for (auto op : context->getRegisteredOperations())
    op.getCanonicalizationPatterns(*patterns, context);
}

// Adds the list of ops that are supported on TPU through constant folding which
// may depend on the inputs shapes not known at this point. Such ops may not
// have any legalization or canonicalization patterns but shouldn't be marked
// for outside compilation.
//