builder.getStrArrayAttr(
                         {"tf_saved_model.session_initializer_restore_op"}));
  init_func->setAttr(
      kTfSavedModelInitializerTypeAttr,
      builder.getStringAttr(kTfSavedModelInitializerRestoreType));

  builder.setInsertionPointToStart(&init_func.front());
  builder.create<func::ReturnOp>(loc, /*operands=*/ValueRange{});

  SymbolTable symbol_table(module_op);

                                        ",", group_annotation);
      }
      for (auto& op : replica.GetBody()) {
        op.setAttr(TF::kParallelExecAnnotation,
                   builder.getStringAttr(group_annotation));
      }
    }
    replicas.push_back(replica);
  }

  return success();
}

// Creates islands per replica from `tf_device.replicate` region and remap

      }
    }
  }

  builder->setInsertionPointAfter(before_op);
  auto launch = builder->create<tf_device::LaunchOp>(
      before_op->getLoc(), builder->getStringAttr(host_device),
      launch_result_types);
  launch.getBody().push_back(launch_block);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(before_op->getLoc(), launch_results);

      if (stb_paddings_attr.getValues<APInt>()[{i, j}].getSExtValue() +
              addtional_pad !=
          bts_crops_attr.getValues<APInt>()[{i, j}].getSExtValue()) {
        op->setAttr("padding", rewriter.getStringAttr("SAME"));
        break;
      }
    }
  }

  // Set dilations
  op->setAttr("dilations", dilations_attr.value());

  if (expand_op) {

      }
    }
  }

  before ? builder->setInsertionPoint(op) : builder->setInsertionPointAfter(op);
  auto launch = builder->create<mlir::tf_device::LaunchOp>(
      op->getLoc(), builder->getStringAttr(host_device), launch_result_types);
  launch.getBody().push_back(launch_block);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<mlir::tf_device::ReturnOp>(op->getLoc(), launch_results);

  assert(while_wrapper && "While op is expected to be wrapped in a IslandOp");

  // Get the dummy constant.
  OpBuilder builder(while_wrapper);
  auto loc = NameLoc::get(
      builder.getStringAttr("chain_control_outputs@" + while_op.getBody()));
  IslandOp const_wrapper = GetDummyConstant(builder, const_type, loc);

  // Get new operand and result types.

  int index = *branch.getValues<int>().begin();
  if (index < 0 || index >= op.num_branches()) index = op.num_branches() - 1;

  auto func = mlir::cast<SymbolRefAttr>(op.getBranches()[index]);
  auto empty = rewriter.getStringAttr("");
  ReplaceTfOpWithNewOp<PartitionedCallOp>(
      rewriter, op, op.getResultTypes(), op.getOperands().drop_front(), func,
      /*config=*/empty, /*config_proto=*/empty, /*executor_type=*/empty);
  return success();

    ArrayAttr layout;  // filled in during the xla-adjust-layout pass
    auto data_and_token =
        rewriter.create<InfeedOp>(op.getLoc(), result_types, token,
                                  /*infeed_config=*/rewriter.getStringAttr(""),
                                  /*layout=*/layout);

    result_types.pop_back();  // remove the token type.

    if (op.get_XlaSharding().has_value()) {

  auto new_concat = builder.create<TFL::ConcatenationOp>(
      getLoc(), getType(), non_empty_values,
      builder.getIntegerAttr(builder.getIntegerType(32), getAxis()),
      builder.getStringAttr(getFusedActivationFunction()));
  return new_concat.getResult();
}

//===----------------------------------------------------------------------===//
// CustomOp