}
    // Use the min/max from the operands and the num_bits and narrow_range
    // 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(

namespace TFR {

namespace {

// Quantize the float value based on given scale and zero point attributes.
IntegerAttr Quantize(float value, Attribute scale_attr, Attribute zp_attr,
                     OpBuilder builder) {
  double scale = mlir::cast<FloatAttr>(scale_attr).getValueAsDouble();
  int64_t zp = mlir::cast<IntegerAttr>(zp_attr).getInt();

  int quantized = static_cast<int>(std::round(value / scale) + zp);
  quantized =

    }
    // Use the min/max from the operands and the num_bits and narrow_range
    // 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(

void HandleConv2DStride(TF::Conv2DOp conv2d) {
  MLIRContext* context = conv2d.getContext();
  SmallVector<int64_t, 4> values = {1, 1, 1, 1};
  auto attrs = llvm::map_range(values, [context](int64_t v) -> Attribute {
    return IntegerAttr::get(IntegerType::get(context, 64), v);
  });
  // TODO(b/157276506): change type of strides to DenseElementsAttr
  auto strides = ArrayAttr::get(context, llvm::to_vector<4>(attrs));
  conv2d->setAttr("strides", strides);

          window_dimensions.getValues<IntegerAttr>()[i].getInt());
    }

    for (auto i = 0; i < window_strides.size(); ++i) {
      window_strides_vec.push_back(
          window_strides.getValues<IntegerAttr>()[i].getInt());
    }

    for (auto i = 0; i < base_dilations.size(); ++i) {
      base_dilations_vec.push_back(
          base_dilations.getValues<IntegerAttr>()[i].getInt());
    }

    if (alias_info.input_index == kUnassigned ||
        alias_info.output_index == kUnassigned)
      continue;
    auto aliasing_attr = device_func.getArgAttrOfType<mlir::IntegerAttr>(
        alias_info.input_index, kAliasingAttr);

    // Set only if aliasing attribute does not exist.
    if (!aliasing_attr) {
      device_func.setArgAttr(
          alias_info.input_index, kAliasingAttr,

using OptionalOpShardingVector =
    llvm::SmallVector<std::optional<OpShardingVariant>, 8>;
using llvm::StringRef;
using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::Builder;
using mlir::IntegerAttr;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OpOperand;
using mlir::OpResult;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::Value;

    cost_map_[op] = iter->second(context, op);
    return;
  }

  // Try to use the recorded cost if any.
  if (cost_recorder_ != nullptr) {
    const auto op_key_attr =
        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
    if (op_key_attr) {
      cost_map_[op] = cost_recorder_->GetCost(op_key_attr.getInt());
      return;
    }
  }

  // These ops are cheap regardless of their input sizes.
  //

  llvm::SmallSet<int, 4> dynamic_arg_idx_set;
  if (dynamic_arg_idx) {
    for (auto idx : dynamic_arg_idx.getValue()) {
      dynamic_arg_idx_set.insert(mlir::dyn_cast<IntegerAttr>(idx).getInt());
    }
  }

  for (auto operand_type_and_idx : llvm::enumerate(op.getOperandTypes())) {
    Type operand_type = operand_type_and_idx.value();
    int index = operand_type_and_idx.index();

  if (elem_type.isSignlessInteger() || elem_type.isUnsignedInteger()) {
    auto start_attr = start_tensor.getValues<IntegerAttr>()[0];
    auto limit_attr = limit_tensor.getValues<IntegerAttr>()[0];
    auto delta_attr = delta_tensor.getValues<IntegerAttr>()[0];
    int num_elements;
    if (elem_type.isUnsignedInteger()) {
      uint64_t start = start_attr.getUInt();