// Gather op patterns.
//===----------------------------------------------------------------------===//

class GetI32Attr<int x>: NativeCodeCall<
  "$_builder.getI32IntegerAttr(" # x # ")">;

class GetI64Attr<int x>: NativeCodeCall<
  "$_builder.getI64IntegerAttr(" # x # ")">;

// The deprecated validate_indices attribute is NOT verified during
// canonicalization because it is neither supported in the old bridge

    "$_builder.create<TF::ReadVariableOp>("
    "$0.getLoc(), GetResourceSubtype($1), $1)">;

def CreateTFSelectOp: NativeCodeCall<
    "$_builder.create<TF::SelectOp>("
    "$0.getLoc(), $3.getType(), $1, $2, $3)">;

def ConstAttrIfThenElse: NativeCodeCall<
    "($0).getValue() ? ($1) : ($2)">;

// Convert clamp(lo, x, hi) to clipbyvalue(x, lo, hi).
def Clamp: NativeCodeCall<
    "$_builder.create<TF::ClipByValueOp>("

include "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/IR/PatternBase.td"

def CreateNoneValue : NativeCodeCall<
  "$_builder.create<TFL::NoValueOp>($0.getLoc(), $_builder.getUnitAttr())">;

// Returns shape of a ranked tensor.
// if called without a ranked tensor it will fail.
def GetShape: NativeCodeCall<"GetShape($0)">;

include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// TODO(jpienaar): Move this somewhere general.
class GetI64ScalarElementsAttr<int value> :
  NativeCodeCall<"DenseElementsAttr::get<int64_t>(RankedTensorType::get({}, $_builder.getIntegerType(64)), " # value # ")">;

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op
     (TF_MapDatasetOp $input_dataset, $other_arguments, $f, $output_types,

  if (::tflite::optimize::QuantizeWeights(
          &q_builder, input_model, quantized_type, use_updated_hybrid_scheme,
          ::tflite::optimize::QuantizerType::OLD_QUANTIZER) != kTfLiteOk) {
    return absl::InvalidArgumentError(
        "Quantize weights transformation failed.");
  }
  const uint8_t* q_buffer = q_builder.GetBufferPointer();
  *result =

    return absl::OkStatus();
  }

  Status LegalizeSingleOp(Operation& op) {
    SourceMgrDiagnosticHandler sourceMgrHandler(source_manager_, &context_);

    OpBuilder op_builder(&op);
    EmptyPatternRewriter pattern_rewriter(op_builder);

    LogicalResult result =
        Tf2XlaRewriter::RewriteOp(&op, pattern_rewriter,
                                  /*device_type=*/"XLA_CPU_JIT");
    if (!result.succeeded()) {

include "mlir/Dialect/Arith/IR/ArithOps.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.td"

class Quantize<string value> : NativeCodeCall<"TFR::Quantize(" # value # ", $0, $1, $_builder)">;

class HasStringAttr<string value> : AttrConstraint<
    CPred<"$_self.cast<StringAttr>().getValue() == \"" # value # "\"">>;

def QuantActRangeNonePattern :
  Pattern<
    (TFR_TFRQuantActRangeOp

using QuantMethod = tensorflow::quantization::QuantizationMethod::PresetMethod;

class EmptyPatternRewriter : public mlir::PatternRewriter {
 public:
  explicit EmptyPatternRewriter(const OpBuilder& op_builder)
      : mlir::PatternRewriter(op_builder) {}
  ~EmptyPatternRewriter() override = default;
};

class TfToUniformAttributeUtilsTestPeer {
 public:
  explicit TfToUniformAttributeUtilsTestPeer() = delete;

include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// Marks the op as no fallback.
def MarkNoFallback : NativeCodeCall<"SetNoFallbackAttr($_builder, $0)">;

def NoFallbackAttrNotSet : Constraint<CPred<
  "!$0.getDefiningOp()->hasAttr(kNoFallbackAttr)">>;

class FloatValueEquals<string val> : Constraint<CPred<
  "FloatValueEquals($0, " # val # ")">>;

                    (TFL_HardSwishOp $input)>;

def IsNchwLayoutOp: Constraint<CPred<
  "$0.get(\"is_nchw_op\") && $0.get(\"is_nchw_op\").dyn_cast<BoolAttr>() "
      "== mlir::BoolAttr::get($_builder.getContext(), true)">>;

def IsNhwcLayoutOp: Constraint<Neg<IsNchwLayoutOp.predicate>>;

// TODO(b/343278954): Move the creation of transposes to a separate prepare pass