const XlaShapeLayoutHelpers::ShapeDeterminationFns& shape_determination_fns,
    bool use_tuple_args, XlaCompiler::CompilationResult* compilation_result,
    std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
    const std::vector<TensorShape>& arg_shapes,
    std::vector<tpu::ShardingAndIndex>* arg_core_mapping,
    std::vector<std::vector<xla::Shape>>* per_core_arg_shapes);

    rollout_config.AllowForDeviceInXlaCompileOnDemand(DeviceType(device_type));

    auto device_factory = DeviceFactory::GetFactory(device_type);
    SessionOptions options;
    std::vector<std::unique_ptr<Device>> devices;
    Status s = device_factory->CreateDevices(
        options, "/job:worker/replica:0/task:0", &devices);
    device_ = std::move(devices[0]);

    tensorflow::AllocatorAttributes host_alloc_attr;

/// threads under a coordinator to stop. This function could be called by any
/// thread or any client.
/// Usage, in the client:
///   Coordinator coord;
///   std::unique_ptr<QueueRunner> qr(&coord, ...);
///   qr.Start(session);
///   coord.RegisterRunner(std::move(qr));
///   /// do some work
///   TF_CHECK_OK(coord.Join());

  char* buf = nullptr;
  Py_ssize_t length;
  std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);

  if (tflite::python_utils::ConvertFromPyString(data, &buf, &length) == -1) {
    PyErr_Format(PyExc_ValueError, "Failed to convert input PyObject");
    return nullptr;
  }
  std::unique_ptr<tflite::FlatBufferModel> model =
      tflite::FlatBufferModel::BuildFromBuffer(buf, length,

  NoOpPass() = default;

  llvm::StringRef getArgument() const final { return "no-op-pass"; }

  void runOnOperation() override {
    // Noop pass does nothing on the operation.
  }
};

std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> CreateNoOpPass() {
  return std::make_unique<NoOpPass>();
}

class ParentPass
    : public mlir::PassWrapper<ParentPass,

    : public DeviceCompilerClient<xla::PjRtLoadedExecutable, xla::PjRtClient> {
 public:
  explicit PjRtDeviceCompilerClient(xla::PjRtClient* client)
      : client_(client) {}

  absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> BuildExecutable(
      const XlaCompiler::Options& options,
      const XlaCompiler::CompilationResult& result) override;

  // Returns a platform-specific serialization of `executable`. The

  server_def.set_job_name("worker");
  server_def.set_task_index(0);
  std::unique_ptr<tensorflow::ServerInterface> w0;
  ASSERT_TRUE(
      factory->NewServer(server_def, ServerFactory::Options(), &w0).ok());
  ASSERT_TRUE(w0->Start().ok());
  server_def.set_task_index(1);
  std::unique_ptr<tensorflow::ServerInterface> w1;
  ASSERT_TRUE(

};

// Struct that wraps TF_Tensor and TF_Status to delete once out of scope.
using Safe_TF_TensorPtr = std::unique_ptr<TF_Tensor, TFTensorDeleter>;
using Safe_TF_StatusPtr = std::unique_ptr<TF_Status, TFStatusDeleter>;

// Used to pass the operation node name from kernel construction to
// kernel computation.
struct HistogramSummaryOp {
  std::string op_node_name;
};

  Tape tape(/*persistent=*/false);
  tape.Watch(inputs[0]);  // Watch x.
  AbstractTensorHandle* exp_output;
  TF_RETURN_IF_ERROR(ops::Exp(ctx, inputs[0], &exp_output, "Exp"));
  std::unique_ptr<GradientFunction> gradient_function(
      new PassThroughGradientFunction);
  tape.RecordOperation(inputs, {exp_output}, gradient_function.release());
  TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx,

                              const SavedObjectGraph& obj_graph,
                              const PartiallyRevivedObjects& objects,
                              std::unique_ptr<TFConcreteFunction>* out) {
  const auto& capture_node_ids = builder.saved_concrete_func->bound_inputs();
  std::vector<ImmediateExecutionTensorHandle*> captures;
  captures.reserve(capture_node_ids.size());