let summary = "Transform Einsum to other TF Ops for the supported variants";
  let constructor = "TF::CreateTransformEinsumPass()";
}

def LegalizeTFGToTFPass : Pass<"tfe-legalize-tfg", "ModuleOp"> {
  let summary = "Legalize from TFG to the TFE dialect";
  let constructor = "TF::CreateLegalizeTFGToTFEPass()";
}

def ReplicateToIslandPass : Pass<"tf-replicate-to-island", "mlir::func::FuncOp"> {

#include "tensorflow/core/lib/math/math_util.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/util/padding.h"

namespace mlir {
namespace odml {
namespace {

#define DEBUG_TYPE "tf-legalize-hlo"

#define GEN_PASS_DEF_LEGALIZEHLOTOTFPASS
#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h.inc"

class LegalizeHloToTf : public impl::LegalizeHloToTfPassBase<LegalizeHloToTf> {

   (HasOneUse $relu_out),
   (HasOneUse $mul_out),
   (HasOneUse $input_neg_out)]>;

// The constant folding in this pass might produce constant in the tf dialect.
// This rule is to legalize these constant to the tfl dialect.
def LegalizeConstOp : Pat<
  (TF_ConstOp ElementsAttr:$value), (TFL_ConstOp $value)>;

// Reorders adds to allow constant folding.
// Add --> Add $input, $constantA

    "/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;

// Run legalize pass and then optimize pass, and make sure some fusing is applied.
// RUN: tf-opt %s -tfl-legalize-tf -tfl-optimize | FileCheck --check-prefix=Fusing %s
// Run legalize pass and then optimize pass, and make sure some fusing is applied, but no mul->fc.
// RUN: tf-opt %s -tfl-legalize-tf -tfl-optimize='disable-fuse-mul-and-fc=true' | FileCheck --check-prefix=NoFusing %s

                               : OpBuilder::atBlockEnd(body, listener);
    }
  }];

  // Canonicalizer wasn't defined for this one. In practise, we legalize the
  // tf.IfOp to scf.If op first and then legalize it to tfl.if to reduce
  // code redundancy.
}

def TFL_WhileOp : Op<TFL_Dialect, "while", [
    DeclareOpInterfaceMethods<LoopLikeOpInterface, ["isDefinedOutsideOfLoop"]>,

namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_LEGALIZEVARIABLESPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

// Attribute name to identify whether variables should be legalized to TFLite or
// not.
const char kLegalizeTflVariables[] = "tfl._legalize_tfl_variables";

bool HasSupportedElementType(Operation* op) {
  return utils::IsSupportedVariableType(op);
}

//   . 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.

After serialization, this pass removes the `_entry_function` attribute from
`tf.XlaCallModule`, and removes all the serialized stablehlo functions
from the top-level module.
### `-tfe-legalize-tfg`

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

namespace mlir {
namespace mhlo {

// Given the Operation, return True if the op is legalized with MLIR.
// Returns false if it is legalized with an old TF2XLA Kernel.
bool IsOpLegalizedWithMlir(Operation& op);

// Given the type ID, check if it's legalized with MLIR.
bool IsTypeLegalizedWithMlir(const TypeID& type_id);

// Returns true if the op is considered a dynamic padder op.