&arg_core_mapping, &per_core_arg_shapes, client)
                     .status());
  }
};

TEST_F(LegalizeTFQuantTest, LegalizesModuleWithTFUniformQuantization) {
  constexpr char mlir_module_string[] = R"mlir(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main(%arg0 : tensor<1xf32>) -> tensor<1xf32> {

#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace quant {
namespace {

// Prefix and suffix to the QuantizationUnit string representation.
constexpr std::string_view kQuantizationUnitPrefix = "QuantizationUnit(";
constexpr std::string_view kQuantizationUnitSuffix = ")";

// Concatenates node name and func name with a "@" separator.
std::string ConcatNodeAndFuncName(std::string_view node_name,

namespace mlir {
namespace TFL {

namespace {

// TODO(b/162842801): Consolidate all util definitions of kTFImplements.
constexpr char kTFImplements[] = "tf._implements";
constexpr char kCustomSSDPostprocessing[] = "TFLite_Detection_PostProcess";
constexpr char kTfNMSPadded[] = "non_max_suppression_padded_v2";

inline ConstBytesAttr CustomOption(OpBuilder* builder,

 public:
  // Number of input arguments for the compiled computation.
  static constexpr size_t kNumArgs = {{ARG_NUM}};

  static constexpr size_t kNumResults = {{RESULT_NUM}};

  // Number of variables for the compiled computation.
  static constexpr size_t kNumVariables = {{VARIABLE_NUM}};

  // Byte size of each argument buffer. There are kNumArgs entries.

using ParsedName = ::tensorflow::DeviceNameUtils::ParsedName;

constexpr const char *kHostAttr = "host";
constexpr const char *kDeviceAttr = "device";
constexpr const char *kTFDeviceAttr = "tf.device";
// TODO(donglin): Handle the case where the address of localhost is different
// from /job:localhost/replica:0/task:0.
constexpr const char *kLocalhost = "/job:localhost/replica:0/task:0";
constexpr const char *kErrorMessage =

  // tf.LookupTableFindV2 ops are usually more costly than tf.AddV2 with the
  // same input size, as it involves more operations like hashing, map lookup,
  // etc.
  constexpr int64_t kLookupTableFindCostScale = 8;
  constexpr int64_t kLookupTableFindStringKeyCostScale = 16;

  auto value_type = mlir::cast<mlir::TensorType>(op.getValues().getType());
  auto key_type = mlir::cast<mlir::TensorType>(op.getKeys().getType());

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace TFDevice {

namespace {

constexpr char kReplicationAttr[] = "mhlo.is_same_data_across_replicas";
constexpr char kMirroredVariableIndicesAttr[] = "_mirrored_variable_indices";

#define GEN_PASS_DEF_ANNOTATEPARAMETERREPLICATIONPASS

    pm->addPass(mlir::mhlo::createLegalizeTFPass(
        /* legalize_chlo=*/true, llvm::StringRef("gpu/xpu"),
        /* prefer_tf2xla=*/false));
  };
}

TEST(XlaLegalizeTest, IllegalOp) {
  constexpr char kMlirIllegalOpStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> tensor<1xi32> {
      %0 = "tf.DoesntExist"() : () -> tensor<1xi32>

  // `auto_control_deps.py` where it is assumed that allocated resource values
  // NEVER alias. We should align our assumptions in the future.
  static constexpr int64_t kUnknownResourceId = -1;
  static constexpr int64_t kInvalidResourceId = -2;
  static constexpr int64_t kMaxResourceTypeId = 9999;
};

}  // namespace detail

// An analysis that runs on a module and maps each resource-type value to a

namespace tensorflow {

using ::testing::_;
using ::testing::NiceMock;
using ::testing::Return;
using ::testing::Test;

constexpr char kOk[] = "OK";
constexpr char kInvalidArgument[] = "INVALID_ARGUMENT";
constexpr char kSuccess[] = "kSuccess";
constexpr char kFailure[] = "kFailure";

class MockMlirOptimizationPass : public MlirOptimizationPass {
 public: