/// `prefer_tf2xla` means an op will be included iff it is not in
/// `MlirLegalizedUnderPreferTf2XlaSet`. `!prefer_tf2xla` mean an op will be
/// included if there is no native MLIR legalization for the op.
void PopulateLegalizeTfWithTf2XlaPatterns(llvm::StringRef device_type,
                                          RewritePatternSet& patterns,
                                          MLIRContext* ctx,

  // Tries to legalize the specified TensorFlow op, if supported.
  //
  // Emits an error and returns failure if an error is encountered during
  // conversion. Note that success return value doesn't mean successful
  // legalization.
  mlir::LogicalResult LegalizeOp();

  // Converts the given operand to expression of kind kConstant or kXlaOp.
  // Emits a remark and returns expression of kind kInvalid on failure.

// convert TF uniform quantized ops to the corresponding quantized MHLO ops.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertTFQuantOpsToMHLOPass();

// TODO(b/288094093): Migrate uniform quantization legalization in a separate
// pass.
void PopulateLegalizeTfQuantizationPatterns(MLIRContext *context,
                                            RewritePatternSet *patterns);

namespace internal {

// Compiles a serialized MLIR module and returns a serialized MLIR module of the
// result of running all the MLIR Bridge passes. If compile_to_xla_hlo is true
// then those passes include all the Legalization to XLA HLO which is returned
// in the compilation_result.
absl::StatusOr<std::string> CompileFromMlirToXlaHlo(
    bool lower_to_xla_hlo, const tpu::MlirToHloArgs& computation,

  EXPECT_TRUE(status);
  EXPECT_EQ(legalize_failure_count.Read("tf.DoesntExist", "Unknown"), 1);
}

TEST(XlaLegalizeTest, LegalOp) {
  // We expect legalization to fail for legal op with dynamic shapes:
  static constexpr char kMlirLegalOpStr[] = R"(
   func.func @infeed_dequeue_tuple_dynamic_error() -> (tensor<3x3xf32>, tensor<4x?xf32>) {

// This test runs twice:
//   1. Through FILECHECK_OPTS="" FileCheck with chlo legalization disabled since verifying
//      that the chlo ops emit produces more useful tests.
//   2. With chlo legalization enabled, verifying diagnostics to pick up any
//      issues with the full lowering (can catch some broadcasting corner
//      cases which emit with a warning).

  if (failed(applyPartialConversion(getOperation(), target,
                                    std::move(patterns)))) {
    getOperation().emitError("mhlo to TFLite legalization failed.");
    signalPassFailure();
  }
}
}  // namespace

void PopulateLegalizeHloToTFLitePatterns(RewritePatternSet* patterns,
                                         MLIRContext* context) {

// is not supported in TensorFlow because TensorFlow requires filters in the
// HWIO data format.
//
// Motivation to prepare for the TFLite legalization before the actual
// legalization is to exploit constant folding opportunities in any newly
// created ops by leveraging constant folding support for the TensorFlow ops.
// This way TFLite can be used as a serialization format only and does not

    if (IsSupportedTFDataForwardingOp(op)) {
      return WalkResult::advance();
    }

    // If any of the operands is a resource type, then we break
    // and mark the module as not valid for TFLite legalization.
    // Note: this might disable native variables in more than needed cases.
    // TODO(b/189370197): Enhance variable analysis.
    for (auto operand : op->getOperands()) {
      if (mlir::isa<TF::ResourceType>(

    // Currently for dot, PTQ supports per-tensor quantization.
    if (rhs_quantized_dimension != -1) {
      return rewriter.notifyMatchFailure(
          op, "Legalization supports only rhs_quantization_axis -1.");
    }
    auto rhs_type = GetUniformQuantizedType(
        op, op.getRhs().getType(), op.getRhsScales(), op.getRhsZeroPoints(),