using ::tsl::Status;
using ::tsl::StatusOr;
using ::xla::ReplicaGroup;
using ::xla::ShapeUtil;
using ::xla::XlaBuilder;
using ::xla::XlaComputation;
using ::xla::XlaOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

#include "tsl/platform/status.h"
#include "tsl/platform/statusor.h"

namespace mlir {
namespace mhlo {

using func::FuncOp;
using mlir::ModuleOp;

static constexpr char kMlirModuleStr[] = R"(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1442 : i32}} {

constexpr char kEntryFuncName[] = "main";

static constexpr char kMlirModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() -> () {
    func.return
  }
})";

MlirToHloArgs CreateTestMlirToHloArgs(const char* module_str = kMlirModuleStr) {
  MlirToHloArgs mlir_to_hlo_args;
  mlir_to_hlo_args.rollout_state =