// Flags associated with the XLA bridge's mark_for_compilation_pass module.
struct MarkForCompilationPassFlags {
  XlaAutoJitFlag xla_auto_jit_flag;

  // Minimum number of operators in an XLA compilation. Ignored for operators
  // placed on an XLA device or operators explicitly marked for compilation.
  int32 tf_xla_min_cluster_size;

  // Maximum number of operators in an XLA compilation.
  int32 tf_xla_max_cluster_size;

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/local_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_future.h"
#include "xla/pjrt/pjrt_stream_executor_client.h"
#include "xla/pjrt/tracked_device_buffer.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "xla/stream_executor/platform_manager.h"
#include "tensorflow/core/common_runtime/dma_helper.h"

  TF_ASSIGN_OR_RETURN(xla::Shape shape,
                      shape_determination_fns.shape_representation_fn(
                          cpu_tensor->shape(), cpu_tensor->dtype(),
                          /*fast_mem=*/false, layout_preference));
  const xla::Layout* device_layout = &(shape.layout());
  // The device id should match the local_hardware_id in
  // tensorflow/compiler/xla/pjrt/pjrt_client.h.
  TF_ASSIGN_OR_RETURN(

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

XlaDeviceCompiler* CreateXlaDeviceCompiler(

};

using XlaExecutableClosure =
    ExecutableClosure<xla::LocalExecutable, xla::LocalClient>;
using XlaExecutableClosureStore =
    ExecutableClosureStore<xla::LocalExecutable, xla::LocalClient>;
using PjRtExecutableClosure =
    ExecutableClosure<xla::PjRtLoadedExecutable, xla::PjRtClient>;
using PjRtExecutableClosureStore =
    ExecutableClosureStore<xla::PjRtLoadedExecutable, xla::PjRtClient>;

#include <utility>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/pjrt_c_api_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"

#include "tensorflow/compiler/tf2xla/type_util.h"
#include "xla/error_spec.h"
#include "xla/literal.h"
#include "xla/literal_util.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/pjrt/cpu/cpu_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_executable.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/tests/literal_test_util.h"
#include "tensorflow/core/framework/tensor.h"

#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/xla_computation.h"
#include "xla/hlo/ir/hlo_sharding.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
#include "xla/shape.h"
#include "xla/translate/mhlo_to_hlo/layout_util.h"
#include "xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

        "@local_xla//xla/client:padding",
        "@local_xla//xla/client:sharding_builder",
        "@local_xla//xla/mlir_hlo",
        "@local_xla//xla/mlir_hlo:convert_op_folder",
        "@local_xla//xla/mlir_hlo:mhlo_passes",
        "@local_xla//xla/mlir_hlo:type_conversion",
        "@local_xla//xla/stream_executor/tpu:c_api_conversions",
        "@local_xla//xla/stream_executor/tpu:tpu_api",

        "@local_xla//xla:status_macros",
        "@local_xla//xla/client:local_client",
        "@local_xla//xla/pjrt:pjrt_client",
        "@local_xla//xla/pjrt:pjrt_future",
        "@local_xla//xla/pjrt:pjrt_stream_executor_client",
        "@local_xla//xla/pjrt:tracked_device_buffer",
        "@local_xla//xla/service:shaped_buffer",
        "@local_xla//xla/stream_executor:device_memory_allocator",