#include "tensorflow/core/tpu/tpu_defs.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using XlaDeviceExecutablePersistor =
    DeviceExecutablePersistor<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;
using PjRtDeviceExecutablePersistor =

// Instead, it takes the following arguments that are obtained from
// OpKernelContext in the above function.
// - `device`: the device used to compile the function.
// - `rm`: the resource manager for DeviceCompiler to store JIT-compiled XLA
// computation.
// - `flr`: the FunctionLibraryRuntime for the `function`.
Status CompileToPjRtLoadedExecutable(
    const DeviceBase* device, const XlaPlatformInfo& platform_info,

absl::StatusOr<ResourceMgr*> GetResourceMgrForDeviceCompiler(
    const OpKernelContext& ctx, const DeviceType& device_type) {
  // We store information about the JIT-compiled XLA computation in the
  // ResourceMgr. The DeviceCompiler (which contains the DeviceCompilationCache)
  // is stored in the tfrt_global ResourceMgr for TPU and the Device ResourceMgr
  // for CPU/GPU. This is to make sure the DeviceCompiler's lifecycle is

      XlaDeviceExecutablePersistor::Config(), device_type);
  auto compiler_client = std::make_unique<XlaDeviceCompilerClient>(client);
  auto xla_device_compiler =
      new DeviceCompiler<xla::LocalExecutable, xla::LocalClient>(
          std::move(persistor), std::move(compiler_client));
  core::ScopedUnref xla_device_compiler_ref(xla_device_compiler);

  auto profiler = new DeviceCompilationProfiler();

#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/platform/status.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {

using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

Status CompileToPjRtLoadedExecutable(
    const DeviceBase* device, const XlaPlatformInfo& platform_info,
    const NameAttrList& function,

#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/platform/errors.h"

namespace tensorflow {
namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

XlaCompiler::CompileOptions GetCompileOptions(bool for_pjrt = false) {
  XlaCompiler::CompileOptions compile_options;

// which will be compiled and executed using XLA.  The XlaLocalLaunchOp is
// responsible for handling interactions with the TensorFlow executor.
// Once all inputs are present, and their shapes are known, the op can
// use a 'DeviceCompiler' to compile and execute code which is specific
// to the shapes of input Tensors.
// XlaLocalLaunchOp uses xla::LocalClient::Compile() and
// xla::LocalExecutable::Run(), and passes arguments into/out of XLA in device

    absl::Span<const ArgShapeAndDType> input_arg_shape_and_dtype,
    absl::Span<const TensorHandle* const> input_handles,
    CompilerArgSource compiler_arg_source) {
  using XlaDeviceCompiler =
      DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;

  se::Stream* stream = nullptr;
  if (const DeviceBase::AcceleratorDeviceInfo* accelerator_device_info =
          dev->tensorflow_accelerator_device_info()) {

#include "tsl/lib/core/status_test_util.h"
#include "tsl/platform/status.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace {
using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;
using PjRtDeviceExecutablePersistor =
    DeviceExecutablePersistor<xla::PjRtLoadedExecutable, xla::PjRtClient>;

absl::flat_hash_map<int, const Tensor*> GetVariableSnapshots(

    }                                                       \
  } while (0)

namespace tensorflow {

namespace {
using XlaDeviceCompiler =
    DeviceCompiler<xla::LocalExecutable, xla::LocalClient>;
using PjRtDeviceCompiler =
    DeviceCompiler<xla::PjRtLoadedExecutable, xla::PjRtClient>;

auto* xla_launch_counter = monitoring::Counter<1>::New(
    "/tensorflow/core/xla_launch_counter",