string Name() override;

  void* AllocateRaw(size_t alignment, size_t num_bytes) override;
  void DeallocateRaw(void* ptr) override;
  std::optional<AllocatorStats> GetStats() override;
  bool ClearStats() override;

 private:
  // The stream executor of the device.
  se::StreamExecutor* stream_executor_;
};

// Helper class for managing data transfers between host and XLA devices.

  tf_stats.bytes_reservable_limit = se_stats->bytes_reservable_limit;
  tf_stats.largest_free_block_bytes = se_stats->largest_free_block_bytes;
  return tf_stats;
}

bool XlaDeviceAllocator::ClearStats() {
  if (!stream_executor_->SynchronizeAllActivity()) {
    return false;
  }
  return stream_executor_->ClearAllocatorStats();
}

XlaDeviceContext::XlaDeviceContext(