if (auto* op = op_or_val.dyn_cast<mlir::Operation*>()) {
    auto name_from_loc = mlir::GetNameFromLoc(op->getLoc());
    if (!name_from_loc.empty()) return name_from_loc;
    // If the location is none of the expected types, then simply use name
    // generated using the op type.
    return std::string(op->getName().getStringRef());
  }
  auto val = op_or_val.dyn_cast<mlir::Value>();

  func.walk([&](Operation* op) {
    b.setInsertionPoint(op);
    if (auto dq = llvm::dyn_cast<DequantizeOp>(op)) {
      auto dcast = b.create<quantfork::DequantizeCastOp>(
          dq.getLoc(), dq.getOutput().getType(), dq.getInput());
      dq.getOutput().replaceAllUsesWith(dcast);
      dq.erase();
    } else if (auto q = llvm::dyn_cast<QuantizeOp>(op)) {
      auto qcast = b.create<quantfork::QuantizeCastOp>(

  const SavedObjectGraph object_graph = tf_package->GetObjectGraph();
  auto& nodes = object_graph.nodes();
  // Get a map of the concrete functions to their input / output signatures.
  auto& concrete_functions = object_graph.concrete_functions();
  auto& root = nodes.at(0);
  for (auto& child : root.children()) {
    // The child's name describes the name of the edge that connects to the

  int num_threads = 3;

  TensorShape data_shape({1, 1});
  auto a = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
  auto b = Placeholder(root, DT_INT32, Placeholder::Shape(data_shape));
  auto c = BatchMatMul(root, a, b);
  ClientSession session(root);
  std::vector<Tensor> outputs;

  auto inter_op_threadpool =
      absl::make_unique<CustomThreadPoolImpl>(num_threads);

    }

    auto top_k_attr = backend_config.getAs<IntegerAttr>("top_k");
    auto reduction_dim_attr =
        backend_config.getAs<IntegerAttr>("reduction_dim");
    auto recall_target_attr = backend_config.getAs<FloatAttr>("recall_target");
    auto aggregate_to_topk_attr =
        backend_config.getAs<BoolAttr>("aggregate_to_topk");
    auto reduction_input_size_override_attr =

  devices.reserve(device_map.size());

  for (auto it : device_map) {
    auto device_alias = it.getName().strref();
    auto device_list = mlir::cast<ArrayAttr>(it.getValue());
    llvm::SmallVector<StringRef, 4> device_list_for_alias;
    device_list_for_alias.reserve(device_list.size());

    for (auto device : device_list)
      device_list_for_alias.emplace_back(

    // Sink down constants into the functions.
    for (auto extern_value : region_extern_values) {
      if (!matchPattern(extern_value, m_Constant())) {
        extern_values.insert(extern_value);
        continue;
      }
      // Add constant at start of region.
      auto const_builder = OpBuilder::atBlockBegin(&region->front());
      auto const_value = const_builder.clone(*extern_value.getDefiningOp());

    PatternRewriter& rewriter) {
  int64_t shape_rank = shape.size();
  auto shape_spec_type =
      RankedTensorType::get({shape_rank}, rewriter.getIntegerType(64));
  Type resultType = RankedTensorType::get(shape, element_type);
  auto constant_attr = DenseElementsAttr::get(shape_spec_type, shape);
  auto shape_tensor =
      rewriter.create<TF::ConstOp>(loc, shape_spec_type, constant_attr);

    for (auto it : llvm::zip_first(loop_carried_yield_operands, types)) {
      auto value = std::get<0>(it);
      auto type = std::get<1>(it);
      if (value.getType() == type) {
        args.push_back(value);
      } else {
        if (IsCompatibleTypeWithTFLCastOp(value.getType()) &&
            IsCompatibleTypeWithTFLCastOp(type)) {

  const auto mins = dyn_cast<DenseFPElementsAttr>(min);
  const auto maxs = dyn_cast<DenseFPElementsAttr>(max);
  if (mins && maxs) {
    min_value.reserve(mins.getNumElements());
    max_value.reserve(maxs.getNumElements());
    for (auto it = mins.begin(); it != mins.end(); ++it) {
      min_value.push_back(FloatAttr::getValueAsDouble(*it));
    }
    for (auto it = maxs.begin(); it != maxs.end(); ++it) {