// This op is guaranteed to be a constant as ODS checks IsConstTensor.
      // Check if the number of elements meets the requirement.
      int current_num_elements =
          mlir::cast<ShapedType>(call_op.getOperand(idx).getType())
              .getNumElements();
      if (current_num_elements < min_num_elements_for_weights_) {
        call_op.emitRemark("Quantization is skipped for ")

  // known.
  mlir::Type output_type;
  auto input_type = mlir::cast<mlir::TensorType>(src_input.getType());

  if (input_type.hasRank()) {
    if (input_type.getShape()[split_dimension] == mlir::ShapedType::kDynamic) {
      output_type = input_type;
    } else {
      auto shape = llvm::to_vector<4>(input_type.getShape());
      if (shape[split_dimension] % num_split != 0) {
        return mlir::emitError(

  if (!encoding) return false;

  for (int i = 0; i < ranked_ty.getRank(); ++i) {
    if (ranked_ty.isDynamicDim(i) &&
        encoding.getBounds()[i] == ShapedType::kDynamic) {
      return false;
    }
  }
  return true;
}

bool HasSymbolRefAttr(Operation* op) {
  for (const auto& attr : op->getAttrs()) {
    Attribute attr_value = attr.getValue();

      AttrValue attr_val;
      // For type attributes, we only propagate the element type.
      mlir::Type elt_type = attr.getValue();
      if (auto shaped_type = mlir::dyn_cast<mlir::ShapedType>(elt_type)) {
        elt_type = shaped_type.getElementType();
      }
      TF_RETURN_IF_ERROR(
          ConvertAttribute(elt_type, remove_ref_type, &attr_val));
      list->add_type(attr_val.type());

                                                              values->end());
      state.attributes.set(
          const_op.getValueAttrName(),
          DenseFPElementsAttr::get(
              mlir::dyn_cast<ShapedType>(const_op.getValue().getType())
                  .clone(rewriter.getBF16Type()),
              bfloat16_values));
    }

    rewriter.replaceOp(op, rewriter.create(state)->getResults());

    // Connect `expand_op` with the input of `stb_op`.
    expand_op.setOperand(0, stb_op.getInput());
    // Calculate the shape for expand.
    auto input_shape =
        mlir::cast<ShapedType>(stb_op.getInput().getType()).getShape();
    SmallVector<int64_t, 4> expand_shape(input_shape.begin(),
                                         input_shape.end());
    expand_shape.insert(expand_shape.begin() + expand_axis, 1);

  bool getQuantizableOps(arith::ConstantOp op,
                         QuantizationUnits& quantizable_ops) const {
    // Non-float tensors do not need quantization.
    auto type = mlir::dyn_cast<ShapedType>(op.getType());
    if (!type || !type.getElementType().isF32()) return false;

    Value value = op.getResult();

    // Check whether dynamic range quantization can be applied.

    if (!hasUsageFromQuantizableOp(op)) return failure();

    // Check if the weight size is big enough.
    int num_elements_threshold = quant_options_.min_num_elements_for_weights();
    int num_elements = cast<ShapedType>(op.getType()).getNumElements();
    if (num_elements < num_elements_threshold) {
      op->emitRemark("Quantization is skipped because the op has ")
          << num_elements << " elements which is fewer than the threshold("

  ConstantAttr<MHLO_FftTypeAttr, "::mlir::mhlo::FftType::" # enumStr>;

def GetInnerDimFromValue : NativeCodeCall<
  "GetInnerDimFromValue($0.getType().cast<ShapedType>(), &$_builder)">;

def CheckInnerDimStatic
  : Constraint<CPred<"CheckInnerDimStatic($0.getType().cast<ShapedType>(), &$_builder)">>;

def : Pat<(TF_FFTOp:$res $input),
          (MHLO_FftOp $input, MHLO_FftTypeValue<"FFT">, (GetInnerDimFromValue $res)),

    shape_attrs.reserve(op.getNumResults());
    for (Type ty : op.getResultTypes()) {
      shape_attrs.push_back(TF::ShapeAttr::get(rewriter.getContext(),
                                               mlir::cast<ShapedType>(ty)));
    }

    // Clone the `host_func` in the `host_mlir_module` attribute if it exists
    // and use it for `shape_inference_graph` attribute on XlaHostCompute.
    func::FuncOp cloned_func;