bool need_to_set_input_nodes_quantization_params = false;
  for (const BlockArgument arg : func.getArguments()) {
    auto shaped = mlir::dyn_cast<ShapedType>(arg.getType());
    if (shaped && mlir::isa<FloatType>(shaped.getElementType()) &&
        !has_quantize_op(arg)) {
      need_to_set_input_nodes_quantization_params = true;
      break;
    }
  }

  if (!need_to_set_input_nodes_quantization_params) {
    return false;
  }

    // TODO(hinsu): Handle constants that are not scalar booleans.
    auto cond_type = mlir::dyn_cast<RankedTensorType>(cond.getType());
    if (!cond_type || !cond_type.getShape().equals({}) ||
        !cond_type.getElementType().isInteger(/*width=*/1))
      return failure();

    // Identify the branch to inline.
    bool cond_value = (*cond.value_begin<APInt>()).getSExtValue();
    func::FuncOp func = cond_value ? then_func : else_func;

  SmallVector<int64_t, 4> permutedShape = applyPermutation(
      type.getShape(), isInvert ? invertPermutationVector(perm) : perm);
  return RankedTensorType::get(permutedShape, type.getElementType());
}

static RankedTensorType GetInvertPermutedTensorType(RankedTensorType type,
                                                    ArrayRef<int64_t> perm) {

    if (!const_op) return failure();

    // Only support float tensors.
    auto tensor_type = mlir::dyn_cast_or_null<TensorType>(const_op.getType());
    if (!tensor_type || !tensor_type.getElementType().isF32()) {
      return failure();
    }

    return success(
        mlir::isa_and_nonnull<DenseFPElementsAttr>(const_op.getValue()));
  }

  void rewrite(mlir::stablehlo::TransposeOp op,

  bool need_to_set_input_nodes_quantization_params = false;
  for (const BlockArgument arg : func.getArguments()) {
    auto shaped = mlir::dyn_cast<ShapedType>(arg.getType());
    if (shaped && mlir::isa<FloatType>(shaped.getElementType()) &&
        !has_quantize_op(arg)) {
      need_to_set_input_nodes_quantization_params = true;
      break;
    }
  }

  if (!need_to_set_input_nodes_quantization_params) {
    return false;
  }

                         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.
    for (auto& use : value.getUses()) {
      Operation* user = use.getOwner();

      auto arg_type = mlir::cast<RankedTensorType>(arg.getType());
      assert(arg_type.getRank() == 0);
      llvm::ArrayRef<TensorType> underlying_type =
          mlir::cast<TF::ResourceType>(arg_type.getElementType()).getSubtypes();

      // If the arg type already matches the global_tensor type, we don't need
      // to do anything.
      if (!underlying_type.empty() &&

  int64_t output_size = InferTensorSize(context, value_type);

  int64_t cost = kLookupTableFindCostScale * output_size;

  if (mlir::isa<mlir::TF::StringType>(key_type.getElementType()))
    cost *= kLookupTableFindStringKeyCostScale;

  return cost;
}

// The cost function for tf.GatherV2.
int64_t InferGatherV2Cost(const CostContext& context, mlir::TF::GatherV2Op op) {

                split_dimension, num_split));
      }

      shape[split_dimension] = shape[split_dimension] / num_split;
      output_type =
          mlir::RankedTensorType::get(shape, input_type.getElementType());
    }
  } else {
    output_type = input_type;
  }

  // Creates a split op that splits |src_input| along |split_dimension|.
  llvm::SmallVector<mlir::Type, 4> output_types(num_split, output_type);

  // require "raising" later.
  for (const Type type : inst->getResultTypes()) {
    if (const TensorType tensor_type = mlir::dyn_cast<TensorType>(type)) {
      if (mlir::isa<VariantType>(tensor_type.getElementType())) {
        return false;
      }
    }
  }

  // Operations that execute function calls shouldn't be constant folded.
  if (llvm::isa<TF::WhileOp, TF::CaseOp, TF::IfOp, CallOpInterface>(inst)) {