absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
  op.walk([&](mlir::Operation* op) {
    const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
    if (!cost_attr) return;
    const auto op_key_attr =
        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
    if (!op_key_attr) return;
    op_cost_map[op_key_attr.getInt()] = cost_attr.getInt();
  });
  return op_cost_map;
}

    : public impl::TPUPartitionedOpConversionPassBase<
          TPUPartitionedOpConversionPass> {
  void runOnOperation() override;
};

template <typename T>
LogicalResult ReplacePartitionedOp(IntegerAttr num_cores_per_replica, T op) {
  constexpr bool is_input =
      std::is_same_v<std::decay_t<T>, TF::TPUPartitionedInputOp>;
  static_assert(
      is_input || std::is_same_v<std::decay_t<T>, TF::TPUPartitionedOutputOp>,

using namespace mlir;
using namespace mlir::quantfork;

/// Converts a possible primitive, real expressed value attribute to a
/// corresponding storage attribute (typically FloatAttr -> IntegerAttr).
/// quantizedElementType is the QuantizedType that describes the expressed
/// origValue.
/// Returns a converter Attribute or nullptr if conversion is not possible.
static Attribute convertPrimitiveValueAttr(

  int64_t prev_index = -1;
  for (auto arg_index_attr : getVariableArgIndicesAttr()) {
    if (!mlir::isa_and_nonnull<mlir::IntegerAttr>(arg_index_attr)) {
      return emitOpError() << "variable_arg_indices must be an integer";
    }

    int64_t index = mlir::dyn_cast<mlir::IntegerAttr>(arg_index_attr)
                        .getValue()
                        .getSExtValue();
    if (index < 0) {

}

LogicalResult ConvertSSDPostProcessFunc::AddIntAttr(
    func::FuncOp func, DictionaryAttr attrs, const std::string& attribute,
    flexbuffers::Builder* builder) {
  auto int_attr = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get(attribute));
  if (!int_attr) {
    return func.emitError()
           << attribute.c_str() << " attribute is not set or not an integer";
  }
  builder->Int(attribute.c_str(), int_attr.getInt());

class FunctionType(Type):
    def __init__(self, *args, **kwargs) -> None: ...
    def getResults(self) -> list[Type]: ...

class IntegerAttr(Attribute):
    def __init__(self, *args, **kwargs) -> None: ...
    def get(self, arg0: int) -> IntegerAttr: ...

class IntegerType(Type):
    def __init__(self, *args, **kwargs) -> None: ...
    def get(self, arg0: int) -> IntegerType: ...

class Location:

  if (auto ty = mlir::dyn_cast<FloatType>(element_ty)) {
    identity_attr = FloatAttr::get(ty, static_cast<double>(identity));
  } else if (auto ty = mlir::dyn_cast<IntegerType>(element_ty)) {
    identity_attr = IntegerAttr::get(ty, static_cast<int64_t>(identity));
  } else {
    return {};
  }

  // Fold: Op(Operand, Identity) -> Operand.
  if (rhs_attr && is_valid_broadcasting(lhs_type, rhs_type, result_type)) {

    if (mirrored_variable_indices_attr) {
      for (const auto& mirrored_index : mirrored_variable_indices_attr) {
        mirrored_replicate_args.insert(
            mlir::cast<IntegerAttr>(mirrored_index).getInt());
      }
    }
    auto func =
        llvm::cast<func::FuncOp>(m.lookupSymbol(cluster_func.getFunc()));
    for (auto entry : llvm::enumerate(cluster_func.getOperands())) {

        RankedTensorType::get({2}, rewriter.getF32Type()),
        {static_cast<float>(min.getValueAsDouble()),
         static_cast<float>(max.getValueAsDouble())});
    ElementsAttr axis_stats;
    IntegerAttr axis;

    quantfork::StatisticsOp stats_op = rewriter.create<quantfork::StatisticsOp>(
        op->getLoc(), op.getInput(), layer_stats, axis_stats, axis);
    op.getOutput().replaceAllUsesWith(stats_op.getResult());

def GetOffSet : NativeCodeCall<"TFL::GetOffSet($0, $1)">;

// Attribute Constraint that checks if the attribute value is zero.
def ZeroIntAttr
  : AttrConstraint<CPred<"$_self.cast<::mlir::IntegerAttr>().getInt() == 0">>;

// Checks if the value has rank at most 'n'.
class HasRankAtLeast<int n> : Constraint<
    CPred<"$0.getType().cast<ShapedType>().hasRank() && "