EXPECT_EQ(send_from_host_dtypes[0], DT_INT32);
  EXPECT_EQ(send_from_host_dtypes[1], DT_FLOAT);
  // Verify step 4: nodes marked with XLA cluster and outside compilation attr.
  add_node = node_name_image["add"];
  EXPECT_NE(add_node, nullptr);
  EXPECT_TRUE(HasNodeAttr(add_node->def(), "_xla"));
  EXPECT_TRUE(HasNodeAttr(add_node->def(), "_oc"));
  // Verify step 5: control edges added.
  bool has_control_edge_from_recv_at_host = false;

// 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;

// TF-DAG: return %[[identity]]

// XLA-LABEL: func @conv_with_int_per_layer
// XLA-DAG: %[[w0_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}-0.282878935, -0.211567819
// XLA-DAG: %[[b0_float:.*]] = "tf.Const"() <{value = dense<[-0.0192535277, -5.998660e-03]> : tensor<2xf32>}> : () -> tensor<2xf32>
// XLA-DAG: %[[w1_float:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}0.208403707, 0.478067577

#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(

  compile_result.program_shape =
      xla::ShapeUtil::MakeProgramShape(
          {
              xla::ShapeUtil::MakeShape(xla::F32, {1, 2}),
              xla::ShapeUtil::MakeShape(xla::S64, {3, 4}),
              xla::ShapeUtil::MakeShape(xla::F32, {1}),
              xla::ShapeUtil::MakeShape(xla::F32, {1}),
              xla::ShapeUtil::MakeShape(xla::S32, {5}),
          },
          xla::ShapeUtil::MakeTupleShape({

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"