if (!op.getResults().empty()) p << " : " << op.getResults().size();
}

void ExecuteOpWithAllocator::print(OpAsmPrinter &p) {
  ExecuteOpWithAllocator op = *this;
  p << "(" << op.getAllocator() << ") key("
    << op->getAttrOfType<mlir::IntegerAttr>("op_key").getInt() << ") cost("
    << op->getAttrOfType<mlir::IntegerAttr>("_tfrt_cost").getInt()

//
// This is necessary because for XLA devices the underlying TF allocator returns
// dummy tensors.
//
// `stream` parameter is nullable when running on host.
std::shared_ptr<se::DeviceMemoryAllocator> GetAllocator(
    DeviceBase* device, se::Stream* stream,
    const XlaPlatformInfo& platform_info);

}  // namespace tensorflow

#include "tensorflow/core/protobuf/error_codes.pb.h"

namespace tensorflow {
namespace {

class DummyDevice : public DeviceBase {
 public:
  explicit DummyDevice(Env* env) : DeviceBase(env) {}
  Allocator* GetAllocator(AllocatorAttributes /*attr*/) override {
    return cpu_allocator();
  }
};

// Helper for comparing output and expected output
void ExpectSummaryMatches(const Summary& actual, const string& expected_str) {

      for (int i = 0; i < resources.size(); i++) {
        resource_var_ptrs[resources[i]] = variable_infos[i].var()->tensor();
      }

      std::shared_ptr<se::DeviceMemoryAllocator> allocator =
          GetAllocator(ctx->device(), GetStream(ctx), platform_info);
      XlaComputationLaunchContext launch_context =
          GetLaunchContext(platform_info, ctx, client, allocator.get());

      .set_allowed_devices(allowed_devices_)
      .set_intra_op_parallelism_threads(intra_op_parallelism_threads_);
  return xla::ClientLibrary::GetOrCreateLocalClient(options);
}

Allocator* XlaDevice::GetAllocator(AllocatorAttributes attr) {
  mutex_lock lock(mu_);
  return GetAllocatorLocked(attr);
}

Allocator* XlaDevice::GetAllocatorLocked(AllocatorAttributes attr) {
  if (attr.on_host()) {

    host_alloc_attr.set_on_host(true);
    host_allocator_ = device_->GetAllocator(host_alloc_attr);

    // Get the device allocator. This should give us an AsyncValueAllocator.
    AllocatorAttributes device_alloc_attr;
    device_alloc_attr.set_on_host(false);
    device_allocator_ = device_->GetAllocator(device_alloc_attr);

    // Create the DeviceCompiler to help with compiling executables.

  se::Stream* stream =
      ctx->op_device_context() ? ctx->op_device_context()->stream() : nullptr;
  std::shared_ptr<se::DeviceMemoryAllocator> allocator_ptr =
      GetAllocator(ctx->device(), stream, platform_info_);
  se::DeviceMemoryAllocator* allocator = allocator_ptr.get();
  XlaComputationLaunchContext launch_context(
      client, allocator, client->default_device_ordinal(),

    std::optional<std::set<int>> allowed_devices;
  };

  // Creates a new XLA Device.
  XlaDevice(const SessionOptions& session_options, const Options& options);
  ~XlaDevice() override;

  Allocator* GetAllocator(AllocatorAttributes attr) override
      TF_LOCKS_EXCLUDED(mu_);
  void Compute(OpKernel* op_kernel, OpKernelContext* context) override;
  void ComputeAsync(AsyncOpKernel* op_kernel, OpKernelContext* context,

 public:
  explicit FakeDevice(const DeviceAttributes& device_attributes)
      : Device(nullptr, device_attributes) {}

  Status Sync() override { return errors::Unimplemented("FakeDevice::Sync()"); }

  Allocator* GetAllocator(AllocatorAttributes attr) override { return nullptr; }

  static std::unique_ptr<Device> Make(const string& name, const string& type) {
    DeviceAttributes device_attributes;
    device_attributes.set_name(name);

    if (!shape_or.ok()) {
      return op_->emitRemark()
             << "failed to get shape for expression. " << expr.HumanString();
    }

    tensors.emplace_back(
        device_->GetAllocator(tensorflow::AllocatorAttributes()), expr.dtype(),
        shape_or.value());

    tensorflow::Tensor& tensor = tensors.back();
    tensorflow::XlaExpression::AssignExpressionToTensor(expr, &tensor);