TPUPartitionedOpConversionPass> {
  void runOnOperation() override;
};

template <typename T>
LogicalResult ReplacePartitionedOp(IntegerAttr num_cores_per_replica, T op) {
  constexpr bool is_input =
      std::is_same_v<std::decay_t<T>, TF::TPUPartitionedInputOp>;
  static_assert(
      is_input || std::is_same_v<std::decay_t<T>, TF::TPUPartitionedOutputOp>,

namespace mlir {
namespace TFTPU {
namespace {
constexpr char kStepMarkerLocationAttr[] = "step_marker_location";
constexpr char kUseXlaSpmdAttr[] = "use_spmd_for_xla_partitioning";

constexpr char kBadStringArrayElementMsg[] =
    "bad '{0}' attribute at index {1}, not a string";
constexpr char kBadArrayElementMsg[] =
    "bad '{0}' attribute at index {1} with value '{2}': failed to parse to {3}";
constexpr char kBadArrayAttrLengthMsg[] =

  int64_t operands_size = get_size(inst->getOperandTypes());

  constexpr int kSizeFactor = 2;
// TODO(b/233827625): Remove TF_DISABLE_CONSTANT_FOLDING macro.
#ifdef TF_DISABLE_CONSTANT_FOLDING
  constexpr int64_t kResultsSizeThreshold = 0;
#else
  constexpr int64_t kResultsSizeThreshold = (1 << 23);  // 1 MB
#endif
  constexpr int64_t kOperandsSizeThreshold = (1 << 30);  // 128 MB

using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::util::MessageDifferencer;
using ::tsl::testing::IsOk;
using ::tsl::testing::StatusIs;

using LiftAsFunctionCallTest = QuantizationTestBase;

constexpr absl::string_view kModuleLifted = R"mlir(
  module {

    "/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status",
    "Counts how often a successful export from TF Dialect to Executor Dialect "
    "is",
    "status");

constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";

namespace {

void AddTfDialectToExecutorPasses(OpPassManager &pm) {
  pm.addPass(mlir::TF::CreateTFRegionControlFlowToFunctional());

#include "tsl/lib/core/status_test_util.h"

namespace tensorflow {
namespace tf2xla {
namespace v2 {
namespace {

constexpr char kExportStreamzName[] =
    "/tensorflow/core/tf2xla/api/v2/tf_dialect_to_executor_dialect_status";
constexpr char kExportSuccess[] = "success";
constexpr char kExportFailed[] = "failed";

using mlir::DialectRegistry;
using mlir::MLIRContext;
using mlir::ModuleOp;
using mlir::OwningOpRef;

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

namespace mlir {
namespace TFL {

namespace {

constexpr char kNgrams[] = "tftext:Ngrams";
constexpr char kWhitespaceTokenizer[] = "tftext:WhitespaceTokenizer";
constexpr char kCustomSgnnProjection[] = "tftext:custom:SgnnProjection";
constexpr char kTFImplements[] = "tf._implements";

using mlir::TF::FuncAttr;
using mlir::TF::StringType;

#include "tensorflow/compiler/mlir/lite/stablehlo/odml_converter/passes.h.inc"

constexpr float kOne = 1.0;
const float kOneOverRoot2 = kOne / std::sqrt(2);
constexpr float kHalf = kOne / 2.0;
constexpr float kTolerance = kOne / 1000.0;

// Gets the operation that uses the sole result of given operation
// if there is only one.
Operation* GetUserIfOnlyOne(Operation* op) {
  if (op->getNumResults() != 1) return nullptr;

    "/tensorflow/core/tf2xla/api/v2/dynamism_function_counter",
    "Counts how many functions are dynamic", "has_dynamism");

constexpr char kNotDynamicFunctionName[] = "kNotDynamicFunction";
constexpr char kDynamicFunctionName[] = "kDynamicFunction";

class InputMetricsLoweringPass
    : public impl::InputLoweringMetricsPassBase<InputMetricsLoweringPass> {
 public:

      dyn_cast_or_null<IntegerType>(qi8_type.getStorageType())));
}

using IsOpFullyQuantizedTest = QuantizationTestBase;

TEST_F(IsOpFullyQuantizedTest, TrueIfOpFullyQuantized) {
  constexpr absl::string_view kFullyQuantizedAdd = R"mlir(
    func.func @fully_quantized_add(%arg0: tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {