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));
      }

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace TFTPU {

namespace {

// This pass removes Identity/IdentityN ops from the TPU computation and
// reachable functions.
// TODO(lyandy): Remove this pass once resource op lifting is migrated to use
// resource alias analysis and support region based control flow. Removing

  // TODO(b/263528090): Check the condition when this feature is beneficial.
  bool enable_two_input_tensors = 11;

  // Supports TPU model quantization. If the target model for the quantization
  // is already converted for TPU, this flag may be helpful. Note that this
  // feature may be unstable as it is under the experimental stage.
  bool experimental_enable_tpu_model_support = 12;

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

  // We store information about the JIT-compiled XLA computation in the
  // ResourceMgr. The DeviceCompiler (which contains the DeviceCompilationCache)
  // is stored in the tfrt_global ResourceMgr for TPU and the Device ResourceMgr
  // for CPU/GPU. This is to make sure the DeviceCompiler's lifecycle is
  // maintained appropriately.
  ResourceMgr* rm = nullptr;
  if (device_type == DEVICE_TPU) {

};

struct DatasetIterator : ::mlir::SideEffects::Resource::Base<DatasetIterator> {
  StringRef getName() final { return "DatasetIterator"; }
};

// Special resource type to track TPU Embedding specific ops, which must execute
// but do not have side effects with one another or with resource variable ops.
struct TPUEmbedding : ::mlir::SideEffects::Resource::Base<TPUEmbedding> {

         * It has an added advantage since you can you easily switch to different hardware accelerators (cpu, gpu, tpu) as per the task.
         * All you need is a good internet connection and you are all set.
      * Try to build TF from sources by changing CPU optimization flags.
   
   *Please let us know if this helps.*
   

        "//tensorflow/compiler/mlir/quantization/...",
        "//tensorflow/compiler/mlir/tf2xla/transforms/...",
        "//tensorflow/lite/...",
        "//third_party/cloud_tpu/inference_converter/...",  # TPU Inference Converter V1
    ] + internal_visibility_allowlist(),
)

package(
    # copybara:uncomment default_applicable_licenses = ["@stablehlo//:license"],
    default_visibility = [

#include "tensorflow/compiler/mlir/tensorflow/utils/serialize_mlir_module_utils.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/protobuf/tpu/topology.pb.h"
#include "tensorflow/core/util/device_name_utils.h"

// #include <gmock/gmock.h>
// #include <gtest/gtest.h>

namespace tensorflow {
namespace {

  RecvTPUEmbeddingActivationsOp recv_op;
  if (failed(GetOp(region, &recv_op))) return failure();

  SendTPUEmbeddingGradientsOp send_op;
  if (failed(GetOp(region, &send_op))) return failure();

  // No TPU embedding ops.
  if (!recv_op && !send_op) return success();

  Location loc = recv_op ? recv_op.getLoc() : send_op.getLoc();
  StringRef config = recv_op ? recv_op.getConfig() : send_op.getConfig();