}

  void* UnifiedMemoryAllocate(uint64 size) override {
    CHECK(stream_executor_->unified_memory_allocate);
    return stream_executor_->unified_memory_allocate(&device_, size);
  }

  void UnifiedMemoryDeallocate(void* mem) override {
    CHECK(stream_executor_->unified_memory_deallocate);
    stream_executor_->unified_memory_deallocate(&device_, mem);
  }

  absl::optional<AllocatorStats> GetAllocatorStats() override {

  ASSERT_FALSE(allocate_called);
  void* mem = executor->UnifiedMemoryAllocate(8);
  ASSERT_NE(mem, nullptr);
  ASSERT_TRUE(allocate_called);
  ASSERT_FALSE(deallocate_called);
  executor->UnifiedMemoryDeallocate(mem);
  ASSERT_TRUE(deallocate_called);
}

TEST_F(StreamExecutorTest, GetAllocatorStats) {
  se_.get_allocator_stats = [](const SP_Device* const device,