// attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, res_type, min_value, max_value, quant_dim, num_bits,

    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, input_type, min_value, max_value, quant_dim, num_bits,

    This op returns a scalar string tensor as a key for user to look for the loaded array
    and a future containing the restored tensor.
  }];

  let arguments = (ins
    TF_Tensor:$variable,
    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
  );

  let results = (outs
    TF_Tensor:$array_key,
    MlrtFutureType: $tensor_future
  );
}

  bool IsCompositeFunction(TF::PartitionedCallOp call_op) const {
    if (!call_op->hasAttr(kQuantTraitAttrName)) {
      return false;
    }

    const auto f_attr = call_op.getFAttr().dyn_cast<FlatSymbolRefAttr>();
    if (!f_attr || !f_attr.getValue().starts_with("composite_")) {
      return false;
    }

    bool has_quantized_types = false;
    for (Value input : call_op.getArgs()) {

    // attribute to create the quantization parameter for the new quantize op.
    rewriter.setInsertionPointAfter(tf_op.getOperation());
    IntegerAttr num_bits = rewriter.getI64IntegerAttr(tf_op.getNumBits());
    BoolAttr narrow_range = rewriter.getBoolAttr(tf_op.getNarrowRange());
    Type res_type = tf_op.getType();
    TypeAttr qtype = quant::GetQuantizedTypeAttr(
        rewriter, res_type, min_value, max_value, quant_dim, num_bits,

  return true;
}

static bool FloatValueEquals(const Attribute &attr, double value) {
  auto fp_attr = mlir::dyn_cast_or_null<DenseFPElementsAttr>(attr);
  if (!fp_attr) return false;

  if (fp_attr.isSplat()) {
    return fp_attr.getSplatValue<APFloat>().isExactlyValue(value);
  }
  return llvm::all_of(fp_attr.getValues<APFloat>(), [value](const APFloat &f) {
    return f.isExactlyValue(value);
  });
}

      failed(AddFloatAttr(func, attrs, "h_scale", &fbb)) ||
      failed(AddFloatAttr(func, attrs, "w_scale", &fbb)))
    return failure();
  auto use_regular_nms =
      mlir::dyn_cast_or_null<BoolAttr>(attrs.get("use_regular_nms"));
  if (!use_regular_nms) {
    return func.emitError()
           << "use_regular_nms attribute is not set or not a bool";
  }
  fbb.Int("use_regular_nms", use_regular_nms.getValue());

    Other ops executed by TFRT may make use of $tensor_future.
  }];

  let arguments = (ins
    Arg<TF_Tensor, [{The variable tensor to be loaded}], []>:$variable,
    DefaultValuedAttr<BoolAttr, "false">:$used_by_host
  );

  let results = (outs
    TF_StrTensor:$array_key,
    TF_Tensor: $tensor_future
  );

  TF_DerivedOperandTypeListAttr Tin = TF_DerivedOperandTypeListAttr<0>;

        "getElementTypeOrSelf($_self.cast<ElementsAttr>().getType()).isa<FloatType>()">,
        "float constant tensor">;

// Checks if the boolean value is false.
def IsFalseBoolAttr : AttrConstraint<
  CPred<"!$_self.cast<BoolAttr>().getValue()">>;

// Checks if the value has only one user.
def HasOneUse : Constraint<CPred<"$0.hasOneUse()">>;

// Gets the type of a value.
def GetValueType : NativeCodeCall<"$0.getType()">;

// This pass extracts a CPU computation cluster with `_xla_outside_compilation`
// annotation from the head or tail of a TPU cluster.

namespace {

using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::ModuleOp;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::Region;
using mlir::StringAttr;
using mlir::Type;
using mlir::Value;
using mlir::WalkResult;