DeviceSet device_set;
  for (auto& device : devices) device_set.AddDevice(device.get());
  AddDevicesToOp(*module_ref, &device_set);

  auto devices_attr =
      (*module_ref)->getAttrOfType<mlir::DictionaryAttr>("tf.devices");
  ASSERT_NE(devices_attr, nullptr);
  ASSERT_EQ(devices_attr.size(), 3);

  // CPU device added with an empty metadata.

  }
  std::string tfr_raw_text;
  for (const auto& file : files) {
    string fullpath = io::JoinPath(composite_mlir_dir, file);
    if (env->MatchPath(fullpath, io::JoinPath(composite_mlir_dir, "*.mlir"))) {
      std::string text;
      TF_RETURN_IF_ERROR(ReadFileToString(env, fullpath, &text));
      tfr_raw_text.append(text);
    }
  }

  auto ctx = TFRDecomposeContext::GetFromText(tfr_raw_text, mlir_ctx);

    flatbuffers::FlatBufferBuilder* builder,
    std::map<std::string, uint8_t>* hardware_names) {
  uint8_t index = 0;
  for (auto& func : *funcs) {
    func.walk([&hardware_names, &index](mlir::Operation* op) {
      auto device_name = GetDeviceName(op);
      if (!device_name.has_value()) return;

      auto iter = hardware_names->find(*device_name);
      if (iter == hardware_names->end()) {

                                    Region& region) {
  auto parallel_group_attr =
      launch->getAttrOfType<StringAttr>(TF::kParallelExecAnnotation);
  auto result = region.walk([&](Operation* op) -> WalkResult {
    if (op->getDialect() != tf_dialect) return WalkResult::advance();

    if (parallel_group_attr) {
      op->setAttr(TF::kParallelExecAnnotation, parallel_group_attr);
    }
    auto device_attr = op->getAttr(kDeviceAttr);

  if (seen.contains(v)) {
    return failure();  // infinite loop
  }
  seen.insert(v);

  if (auto blockArg = mlir::dyn_cast<BlockArgument>(v)) {
    Operation *op = blockArg.getOwner()->getParentOp();

    // If we're in the first block, then the argument to that block is the
    // one we're looking for.
    if (auto func = cast<func::FuncOp>(op)) {
      if (blockArg.getOwner()->isEntryBlock()) {
        *out = blockArg;

TEST(GetDenseAttrFromTensorProtoAttrTest, Qint8ToUQ8Succeeds) {
  auto context = CreateContext();
  TensorType result_tensor_type = RankedTensorType::get(
      {2, 2}, quant::UniformQuantizedType::get(
                  quant::QuantizationFlags::FlagValue::Signed,
                  IntegerType::get(context.get(), 8),
                  FloatType::getF32(context.get()), 3.0, 2, -128, 127));

  auto dense_attr =

                            const ArrayRef<TF::VarHandleOp> var_handle_ops) {
  auto builder = OpBuilder::atBlockEnd(&save_func.getBody().front());

  SmallVector<std::string> tensor_names = {};
  SmallVector<Value> tensor_values = {};
  for (auto var_handle_op : var_handle_ops) {
    tensor_names.emplace_back(var_handle_op.getSharedName().str());

    auto read_var_op = builder.create<TF::ReadVariableOp>(

    TF_ASSIGN_OR_RETURN(auto func, ExpandNode(ndef, fname));
    TF_RETURN_IF_ERROR(ctx.AddFunctionDef(func));
  }

  // Rewrite the out_op to be the call op. This essentially a deep copy of the
  // orig_op, except the op name.
  auto* new_op = new EagerOperation(&ctx);
  TF_RETURN_IF_ERROR(
      new_op->Reset(fname.c_str(), orig_op->DeviceName().c_str()));
  for (auto input : orig_op->GetInputs()) {

  TF_ASSERT_OK_AND_ASSIGN(auto graph, BuildOpGraphWithOutputShapes());

  auto build_result = BuildHloFromGraph(*graph, /*use_output_shapes=*/false);
  TF_ASSERT_OK(build_result);
}

TEST(CompileMlirUtil, UsesIncorrectOutputShapesWhenPresent) {
  TF_ASSERT_OK_AND_ASSIGN(auto graph, BuildOpGraphWithOutputShapes());

  auto build_result = BuildHloFromGraph(*graph, /*use_output_shapes=*/true);

    auto mirrored_variable_indices_attr =
        replicate->getAttrOfType<ArrayAttr>(kMirroredVariableIndicesAttr);
    llvm::SmallDenseSet<int64_t, 8> mirrored_replicate_args;
    if (mirrored_variable_indices_attr) {
      for (const auto& mirrored_index : mirrored_variable_indices_attr) {
        mirrored_replicate_args.insert(
            mlir::cast<IntegerAttr>(mirrored_index).getInt());
      }
    }
    auto func =