auto input_shapes_attr =
      func_op->getAttrOfType<ArrayAttr>(kTfInputShapesAttr);
  if (!input_shapes_attr) return success();

  if (input_shapes_attr.size() != func_op.getNumArguments()) {
    return func_op->emitError(
               "Number of arguments and 'tf._input_shapes' "
               "attribute size do not match. ")
           << "Num args: " << func_op.getNumArguments()

                                            bool* changed = nullptr) {
  auto res = result->AddResourceDevice(resource, device, changed);
  if (failed(res)) {
    error_reporting_op->emitError()
        << "Conflicting device assignment for resource";
  }
  return res;
}

// Extracts and canonicalizes the device attribute.
inline StringRef GetDeviceAttr(func::FuncOp func, int arg_no) {

    if (impl_iter == func_impls.end()) {
      call_op.emitError(
          "we cannot find corresponding implementation for this call op");
      signalPassFailure();
    }

    for (auto impl : impl_iter->second) {
      auto inference_device_type = GetInferenceDeviceTypeForOp(impl);
      if (!inference_device_type.has_value()) {
        impl.emitError("we cannot find inference device type for this func");

  for (auto func : getOperation().getOps<func::FuncOp>()) {
    auto walk_result = func.walk([&](func::CallOp op) -> WalkResult {
      if (!symbol_table.lookup<func::FuncOp>(op.getCallee()))
        return getOperation().emitError()
               << func.getName() << " is not in the funcs allowlist";
      return WalkResult::advance();
    });
    if (walk_result.wasInterrupted()) return signalPassFailure();
  }
}

}  // namespace

      module_op, /*initializer_type=*/kTfSavedModelInitializerRestoreType);
  if (init_func_op) {
    if (failed(
            applyPatternsAndFoldGreedily(init_func_op, std::move(patterns)))) {
      init_func_op->emitError(
          "Failed to remove variable assignment by const patterns.");
      signalPassFailure();
    }
  } else {
    LOG(INFO) << "Initializer function with type 'restore_op' does not exist. "

  auto module_op = getOperation();

  patterns.add<CastBf16OpsToF32>(ctx);
  populateWithGenerated(patterns);

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    module_op.emitError() << "quant-cast-bf16-ops-to-f32 failed.";
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the Cast BF16 ops to F32 pass.

    auto interface_name = GetInterFaceName(func);
    if (!interface_name.has_value()) continue;

    auto target = GetTargetAnnotation(func);
    if (!target.has_value()) {
      func.emitError("we cannot get hardware info for this function.");
      signalPassFailure();
    }

    float total_cost = GetCostForFunc(&func, *target);
    OpBuilder builder(func);
    UpdateCost(func, total_cost, &builder);

  int32_t container_size =
      get_resource_op->getAttrOfType<mlir::ArrayAttr>("container").size();

  if (!(indices_size == shared_name_size &&
        shared_name_size == container_size)) {
    return get_resource_op->emitError()
           << "length of attribute arrays do not match. indices = "
           << indices_size << ", shared_name = " << shared_name_size
           << ", container = " << container_size;
  }

  return success();
}

// failure.
template <typename OpT>
LogicalResult GetOp(Region* region, OpT* result) {
  *result = {};
  for (auto op : region->getOps<OpT>()) {
    if (*result) return op.emitError("should be unique within a function");
    *result = op;
  }
  return success();
}

LogicalResult RunOnRegion(Region* region) {
  RecvTPUEmbeddingActivationsOp recv_op;

  patterns.add<RemoveTpuOp>(ctx);
  patterns.add<RemoveIdentity>(ctx);

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    module_op.emitError() << "quant-convert-tpu-model-to-cpu pattern "
                             "conversion did not converge.";
    signalPassFailure();
    return;
  }
}

}  // namespace