break;
        }
      }
      if (!found) {
        resultAttrForOp.push_back(builder.getNamedAttr(
            kFuncDeviceAttr,
            StringAttr::get(func.getContext(), it->getSecond())));
      }
    }
  }

  llvm::SmallVector<DictionaryAttr, 8> tmp;
  for (const auto& res : result_attrs) {
    tmp.push_back(builder.getDictionaryAttr(res));
  }
  func.setAllResultAttrs(tmp);
}

  MLIRContext* ctx = builder.getContext();
  // Collect the shapes of the output to fill up the Sout attribute.
  SmallVector<Attribute> shape_attrs;
  for (const Type result_type : output_types) {
    shape_attrs.push_back(
        tf_type::ShapeAttr::get(ctx, mlir::cast<ShapedType>(result_type)));
  }
  auto empty_array_attr = ArrayAttr::get(ctx, {});
  auto platforms = ArrayAttr::get(ctx, {StringAttr::get(ctx, kPlatformCpu)});

          (operations.size() == 2 &&
           dyn_cast<mlir::TF::NoOp>(operations.front()) &&
           operations.back().hasTrait<OpTrait::IsTerminator>())) {
        to_remove.push_back(init_func_op);
      } else {
        to_keep.push_back(sym_ref);
      }
    }

    for (auto func_op : to_remove) rewriter.eraseOp(func_op);

    if (to_keep.empty())
      rewriter.eraseOp(op);
    else

          }))
        result_types.push_back(result.getType());
    }
  }

  IslandOp new_island = OpBuilder(island).create<IslandOp>(
      island.getLoc(), result_types,
      /*control=*/ControlType::get(island.getContext()),
      /*controlInputs=*/island.getOperands());
  new_island.getBody().push_back(new Block);

    plugin_status.reserve(files.size());
    for (int i = 0; i < files.size(); i++)
      plugin_status.push_back(TF_NewStatus());
    result = ops_->paths_exist(filesystem_.get(), translated_names.data(),
                               files.size(), plugin_status.data());
    for (int i = 0; i < files.size(); i++) {
      status->push_back(StatusFromTF_Status(plugin_status[i]));
      TF_DeleteStatus(plugin_status[i]);
    }
  }

    hosts_by_core.reserve(num_replicas);
    for (int replica = 0; replica < num_replicas; ++replica) {
      devices_by_core.push_back(tpu_devices[replica][core].device);
      hosts_by_core.push_back(tpu_devices[replica][core].host);
    }

    device_attrs.push_back(
        builder->getNamedAttr(tensorflow::GetDeviceAliasForLogicalCore(core),
                              builder->getStrArrayAttr(devices_by_core)));

    variable_names.push_back(node_name);
    NodeDef* node_def = name_to_node_map.at(node_name);
    if (node_def->op() == "VarHandleOp") {
      // If this is a resource variable, we have to run the corresponding
      // ReadVariableOp.
      tensor_names.push_back(node_name + "/Read/ReadVariableOp:0");
    } else {
      tensor_names.push_back(node_name + ":0");
    }
  }

      dim_vars.push_back(absl::StrCat("size_t dim", dim));
      dim_sizes += absl::StrCat("[", shape.dimensions(dim), "]");
      indices += absl::StrCat("[dim", dim, "]");
      count *= shape.dimensions(dim);
    }
  }
  rewrites->push_back({"{{I}}", absl::StrCat(i)});
  rewrites->push_back({"{{TYPE}}", type});
  rewrites->push_back({"{{DIM_VARS}}", absl::StrJoin(dim_vars, ", ")});

    for (const auto& output : specs.outputs) {
      terminal_nodes.push_back(std::string(ParseTensorName(output).node()));
    }
    for (const auto& control_output : specs.control_outputs) {
      terminal_nodes.push_back(std::string(control_output));
    }
    for (const auto& input : specs.inputs) {
      terminal_nodes.push_back(input.first);
    }
    TF_RETURN_IF_ERROR(tensorflow::grappler::SetTransitiveFaninGraph(

                pjrt_device->default_memory_space().value_or(nullptr));
        owned_args->push_back(std::move(pjrt_buffer));
        args->push_back(owned_args->back().get());
      }
    } else {
      if (av_tensor->GetBuffer() == nullptr) {
        CHECK_EQ(tensor->NumElements(), 0);  // Crash OK
        continue;
      }
      args->push_back(av_tensor->GetBuffer().get());
    }

    if (!tensor->RefCountIsOne()) {