//   . Legalizes control flow operations.
//   . Decomposes compound resource operations so that the only remaining
//     operations on resource variables are resource reads/writes..
//   . Replaces resource reads/writes with function inputs/outputs and
//     eliminates the use of resource variables.
//   . Legalizes the operations to XLA HLO operations.

class OperationPass;
class Type;

namespace TFL {

// Creates an instance of the TensorFlow Lite dialect LegalizeTF pass.
// When the given run_tfl_runtime_verification value is true, it will check each
// TFL builtin op towards the TFL runtime capability and the incompatible TF ops
// will be left in the graph without getting legalized. If `preserve_assert_op`
// is true, the TF::AssertOp will not be removed.

      TypeID::get<TF::ModOp>(),

      // MatrixDiagPartV3 should use the MLIR implementation due to performance.
      TypeID::get<TF::MatrixDiagPartV3Op>(),

      // Ops that are legalized in the old bridge using MlirXlaOpKernel
      TypeID::get<TF::AbsOp>(),
      TypeID::get<TF::AtanOp>(),
      TypeID::get<TF::AvgPool3DOp>(),
      TypeID::get<TF::BiasAddGradOp>(),
      TypeID::get<TF::CeilOp>(),

  ];
}

def LegalizeWhilePass : Pass<"tfl-legalize-tf-while", "mlir::ModuleOp"> {
  let summary = "Legalize from TensorFlow While to TensorFlow Lite While.";
  let constructor = "CreateLegalizeTFWhilePass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

def LegalizeVariablesPass : Pass<"tfl-legalize-variables-tf", "mlir::ModuleOp"> {
  let summary = "Legalize TensorFlow variables to TensorFlow Lite dialect.";

    "/tensorflow/core/tf2xla/v1/mlir_failed_xla_legalize_tf_pass_count",
    "Counts the failure of legalization of ops", "op_name", "legality");

class LegalizeTF : public impl::LegalizeTFBase<LegalizeTF> {
 public:
  explicit LegalizeTF(bool legalize_chlo,
                      std::optional<StringRef> tf2xla_fallback_device_type,
                      bool prefer_tf2xla) {
    legalize_chlo_ = legalize_chlo;

def : Pat<(MHLO_ConstantOp:$output $value), (TF_ConstOp $value),
          [(TF_Tensor $output)]>;

//===----------------------------------------------------------------------===//
// Binary op patterns.
// Note that these are legalized from chlo.broadcast_* ops, since those are
// semantically compatible with the corresponding TF ops. Depending on
// context, getting to these ops may require some raising.

// This transformation pass converts operations in TensorFlow dialect into
// operations that are legal in the TensorFlow Lite dialect.  Operations that
// can be legalized to TensorFlow Lite dialect with simple replacements are part
// of this pass and other operations that may create extra ops should be part of
// the PrepareTF pass which should be run before this pass.  That way any

                        lower_to_xla_hlo);

  if (lower_to_xla_hlo) {
    // This pass operates on MHLO control flow ops so it should be legalized
    // after the control flow ops are legalized.
    pm.addPass(mlir::mhlo::CreateLegalizeTFCommunicationPass());

    // Everything should be MHLO after this.
    if (!allow_partial_conversion) {
      pm.addNestedPass<mlir::func::FuncOp>(

}

def TFR_TFRQuantRescaleOp : TFR_Op<"quant_rescale", [Pure]> {
  let description = [{
   The `quant_rescale` rescales the elements of the integer tensor by the
   floating-point rescale factor. This op needs to be legalized to the preferred
   operations of the backends.

    Example:

    ```mlir
    %3 = tfr.quant_rescale(%2, %1, %0) :
        (tfr.tensor, tfr.tensor, i64) -> (tfr.tensor)
    ```
  }];

// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT legalize-chlo=false use-tf2xla-fallback=true prefer-tf2xla=true" %s | FileCheck %s
// RUN: tf-opt "-xla-legalize-tf=device-type=XLA_CPU_JIT legalize-chlo=false prefer-tf2xla=true" %s | FileCheck --check-prefix NOFALLBACK %s

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

// CHECK-LABEL: @abs
func.func @abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {