limitations under the License.
==============================================================================*/
#include "tensorflow/c/experimental/stream_executor/stream_executor_test_util.h"

#include "tensorflow/c/experimental/stream_executor/stream_executor.h"

namespace stream_executor {
namespace test_util {

/*** Functions for creating SP_StreamExecutor ***/
void Allocate(const SP_Device* const device, uint64_t size,

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/stream_executor/host/host_platform_id.h"
#include "xla/stream_executor/platform_manager.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"

        "@local_tsl//tsl/platform:statusor",
        "@local_xla//xla/stream_executor/tpu:c_api_conversions",
        "@local_xla//xla/stream_executor/tpu:status_helper",
        "@local_xla//xla/stream_executor/tpu:tpu_api",
        "@local_xla//xla/stream_executor/tpu:tpu_executor_base",
        "@local_xla//xla/stream_executor/tpu:tpu_node_context",
        "@local_xla//xla/stream_executor/tpu:tpu_platform_interface",

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "xla/layout.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/shape.h"
#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {

static FailureOr<std::vector<int64_t>> GetTPUInfeedLayoutFromAPI(
    RankedTensorType t) {

#ifndef TENSORFLOW_COMPILER_JIT_XLA_HOST_SEND_DEVICE_CONTEXT_H_
#define TENSORFLOW_COMPILER_JIT_XLA_HOST_SEND_DEVICE_CONTEXT_H_

#include "xla/shape.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/status_helper.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/stream_executor/tpu/tpu_node_context.h"
#include "xla/stream_executor/tpu/tpu_platform.h"
#include "xla/stream_executor/tpu/tpu_platform_interface.h"
#include "xla/stream_executor/tpu/tpu_stream_interface.h"

#ifndef TENSORFLOW_COMPILER_JIT_XLA_HOST_RECV_DEVICE_CONTEXT_H_
#define TENSORFLOW_COMPILER_JIT_XLA_HOST_RECV_DEVICE_CONTEXT_H_

#include "xla/shape.h"
#include "xla/stream_executor/device_memory.h"
#include "xla/stream_executor/stream.h"
#include "tensorflow/core/framework/device_base.h"
#include "tfrt/concurrency/async_value_ref.h"  // from @tf_runtime

namespace tensorflow {

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/constant_fold.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/convert_tensor.h"
#include "xla/stream_executor/stream_executor.h"
#include "tensorflow/core/framework/logging.h"
#include "tsl/platform/statusor.h"

namespace mlir {
namespace {

}  // anonymous namespace

#include "tensorflow/compiler/mlir/tensorflow/transforms/set_tpu_infeed_layout.h"
#include "xla/layout.h"
#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
#include "xla/shape.h"
#include "xla/stream_executor/tpu/c_api_conversions.h"
#include "xla/stream_executor/tpu/tpu_api.h"
#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {
namespace mhlo {
namespace {

#define GEN_PASS_DEF_INFEEDOPSXLAADJUSTLAYOUT

# Description:
# test for stream_executor
load(
    "//tensorflow:tensorflow.bzl",
    "tf_cc_shared_object",
)

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    licenses = ["notice"],
)

tf_cc_shared_object(
    name = "test_pluggable_device.so",
    srcs = ["test_pluggable_device.cc"],
    visibility = ["//tensorflow/c:__subpackages__"],
    deps = [