// 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);

    @Override
    public DBDef assistCurrentDBDef() {
        // TODO Auto-generated method stub
        return null;
    }

    @Override
    public DataSource assistDataSource() {
        // TODO Auto-generated method stub
        return null;
    }

    @Override
    public DBMetaProvider assistDBMetaProvider() {
        // TODO Auto-generated method stub
        return null;
    }

    @Override

  // maximal sharding then only the specified logical device take the value as
  // the input.
  for (const auto& sharding_attr_and_index : llvm::enumerate(sharding_attrs)) {
    const auto& sharding_attr = sharding_attr_and_index.value();
    const auto input_index = sharding_attr_and_index.index();
    const auto& input_value = cluster_func_inputs[input_index];

    xla::OpSharding sharding;
    if (DecodeShardingAttribute(

    void (*delete_func)(void*), TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
  auto* resource_mgr = cc_ctx->resource_manager();
  tensorflow::core::RefCountPtr<tensorflow::PluginResource>
      tf_plugin_resource_ptr;
  tensorflow::PluginResource* tf_plugin_resource = nullptr;

  auto cc_status = resource_mgr->LookupOrCreate<tensorflow::PluginResource>(

go install -a -n -pgo=auto ./a/a1
stderr 'compile.*-pgoprofile=.*a1.go'

go run -a -n -pgo=auto ./a/a1
stderr 'compile.*-pgoprofile=.*a1.go'

go test -a -n -pgo=auto ./a/a1
stderr 'compile.*-pgoprofile=.*a1.go.*a1_test.go'
stderr 'compile.*-pgoprofile=.*external_test.go'

# go list commands should succeed as usual
go list -pgo=auto ./a/a1

go list -test -pgo=auto ./a/a1

go list -deps -pgo=auto ./a/a1