}

/*static*/ se::DeviceMemoryBase XlaTensor::DeviceMemoryFromTensor(
    const Tensor& tensor) {
  const XlaTensor* xla_tensor = FromTensor(&tensor);
  if (xla_tensor) {
    CHECK(xla_tensor->has_shaped_buffer());
    return xla_tensor->shaped_buffer().root_buffer();
  } else {
    return se::DeviceMemoryBase(const_cast<char*>(tensor.tensor_data().data()),
                                tensor.tensor_data().size());
  }
}

//
// To distinguish between "normal" device tensors and XlaTensors, the raw
// pointer data stored in the TensorBuffer is a tagged pointer.
class XlaTensor {
 public:
  // Downcast from a Tensor to an XlaTensor. Return nullptr if the downcast
  // fails.
  static XlaTensor* FromTensor(const Tensor* tensor);

  // Create a DeviceMemoryBase from a Tensor. The Tensor can be an XlaTensor, in

limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/jit/xla_device_ops.h"

#include <memory>

#include "tensorflow/compiler/jit/xla_tensor.h"

namespace tensorflow {

XlaDeviceDummyOp::XlaDeviceDummyOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

void XlaDeviceDummyOp::Compute(OpKernelContext* ctx) {

#ifndef TENSORFLOW_COMPILER_JIT_XLA_DEVICE_CONTEXT_H_
#define TENSORFLOW_COMPILER_JIT_XLA_DEVICE_CONTEXT_H_

#include <memory>

#include "absl/synchronization/mutex.h"
#include "tensorflow/compiler/jit/xla_tensor.h"
#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/client/global_data.h"
#include "xla/client/local_client.h"

    num_elements *= shape[i];
  }
  std::vector<float> data(num_elements);
  for (size_t j = 0; j < num_elements; ++j) {
    data[j] = j;
  }
  DLManagedTensor dlm_in = {};
  DLTensor* dltensor_in = &dlm_in.dl_tensor;
  dltensor_in->data = data.data();
  dltensor_in->device = {kDLCPU, 0};
  dltensor_in->ndim = static_cast<int32_t>(shape.size());
  dltensor_in->dtype = {kDLFloat, 32, 1};
  dltensor_in->shape = shape.data();

  // context in tensorflow_accelerator_device_info(). Also,
  // tensorflow_accelerator_device_info() == nullptr is used as an IsCPU test.
  // We need XlaCpuDevice to be treated not as CPU because it allocates
  // XlaTensors, not regular Tensors.
  Status status = device->UseAcceleratorDeviceInfo();
  if (!status.ok()) {
    errors::AppendToMessage(&status, "while setting up ", DEVICE_GPU_XLA_JIT);
    return status;
  }