MLIRContext context;
  auto id =
      StringAttr::get(&context, "//tensorflow/python/keras/keras_file.py");
  auto loc = FileLineColLoc::get(&context, id, 0, 0);
  auto id2 =
      StringAttr::get(&context, "//tensorflow/python/something/my_test.py");
  auto loc2 = FileLineColLoc::get(&context, id2, 0, 0);
  auto id3 = StringAttr::get(&context, "python/tensorflow/show_file.py");
  auto loc3 = FileLineColLoc::get(&context, id3, 0, 0);

    return rewriter.notifyMatchFailure(op, "zero_points must be constant");
  }

  auto original_element_type = getElementTypeOrSelf(original_type);
  if (!mlir::isa<TF::Qint8Type, TF::Qint32Type>(original_element_type)) {
    return rewriter.notifyMatchFailure(
        op, "Quantized type must be qint8 or qint32.");
  }
  auto storage_type = GetIntTypeFromTFQint(original_element_type);

  // has the property that if there is an edge from SCC1->SCC2, SCC1 is visited
  // after SCC2, i.e., the graph is traversed bottom up just the way we want.
  auto scc_begin = llvm::scc_begin(&call_graph);
  auto scc_end = llvm::scc_end(&call_graph);
  for (auto& scc : make_range(scc_begin, scc_end)) {
    // Each SCC node is a collection of callgraph nodes that form a cycle. We

    TFL::QConstOp rhs_qconst_op;

    auto GetBroadcastedConstOp = [&](Value operand) -> TFL::QConstOp {
      if (auto broadcast_op = dyn_cast_or_null<stablehlo::BroadcastInDimOp>(
              operand.getDefiningOp())) {
        auto stablehlo_const_op = dyn_cast_or_null<stablehlo::ConstantOp>(
            broadcast_op.getOperand().getDefiningOp());
        auto const_uniform_quantized_type =

  // Remove return values forwarded from unused arguments.
  auto return_op = func_op.front().getTerminator();
  auto old_return_vals = llvm::to_vector<8>(return_op->getOperands());
  int64_t skipped_retvals = 0;
  for (auto entry : llvm::enumerate(old_return_vals)) {
    auto return_val = entry.value();
    if (auto arg = mlir::dyn_cast<BlockArgument>(return_val)) {
      auto it = infos.find(arg.getArgNumber());

  }

  auto type_from_array_attr = [&node, &idx, &builder](
                                  absl::string_view output_shape_attr,
                                  absl::string_view element_type_attr) {
    auto* output_shapes = node.attrs().Find(output_shape_attr);
    auto* element_types = node.attrs().Find(element_type_attr);
    const auto& output_shape = output_shapes->list().shape(idx);

void RaiseTargetSubgraphsPass::runOnOperation() {
  ModuleOp module = getOperation();
  auto& side_effect_analysis = getAnalysis<TF::SideEffectAnalysis>();
  SmallVector<func::FuncOp> funcs(module.getOps<func::FuncOp>());
  int func_count = -1;
  for (auto func : funcs) {
    const auto& info = side_effect_analysis.GetAnalysisForFunc(func);
    for (auto& block : func) {
      OpBuilder builder = OpBuilder::atBlockBegin(&block);

  TF_RETURN_IF_ERROR(ctx->RegisterFunction(trace.get()));
  auto cleanup = absl::MakeCleanup(
      [fname, ctx]() { ctx->RemoveFunction(fname).IgnoreError(); });
  auto call_op = AbstractOperationPtr(ctx->CreateOperation());
  TF_RETURN_IF_ERROR(
      call_op->Reset(fname.c_str(), /*raw_device_name=*/nullptr));
  for (auto t : inputs) {
    TF_RETURN_IF_ERROR(call_op->AddInput(t));
  }

    auto result_value = gather_nd_first.getResult();
    auto params_value = gather_nd_first.getParams();
    auto indices_value = gather_nd_first.getIndices();

    auto result_type = dyn_cast<RankedTensorType>(result_value.getType());
    auto params_type = dyn_cast<RankedTensorType>(params_value.getType());
    auto indices_type = dyn_cast<RankedTensorType>(indices_value.getType());

using ::testing::NotNull;
using ::tsl::testing::StatusIs;

PJRT_Buffer* CreateCBuffer() {
  auto status = pjrt::PjrtApi(DEVICE_CPU);
  if (!status.ok()) {
    CHECK_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
  }
  auto pjrt_client = xla::GetCApiClient(DEVICE_CPU);
  CHECK_OK(pjrt_client.status());
  auto c_api_client = down_cast<xla::PjRtCApiClient*>(pjrt_client->get());
  std::vector<int32_t> data(1, 0);