Searched refs:host_memory_allocator (Results 1 – 3 of 3) sorted by relevance
199 std::unique_ptr<tensorflow::Allocator> host_memory_allocator, in PjRtStreamExecutorClient() argument205 host_memory_allocator_(std::move(host_memory_allocator)), in PjRtStreamExecutorClient()742 void* staging_buffer = host_memory_allocator()->AllocateRaw( in BufferFromHostBuffer()749 on_delete_callback = [staging_buffer, host_memory_allocator = in BufferFromHostBuffer()750 host_memory_allocator()]() { in BufferFromHostBuffer()751 host_memory_allocator->DeallocateRaw(staging_buffer); in BufferFromHostBuffer()782 void* ptr = host_memory_allocator()->AllocateRaw( in BufferFromHostBuffer()785 ptr, [host_memory_allocator = host_memory_allocator()](void* ptr) { in BufferFromHostBuffer()786 host_memory_allocator->DeallocateRaw(ptr); in BufferFromHostBuffer()
494 auto host_memory_allocator = in GetGpuClient() local510 std::move(host_memory_allocator), in GetGpuClient()
131 std::unique_ptr<tensorflow::Allocator> host_memory_allocator,254 tensorflow::Allocator* host_memory_allocator() const { in host_memory_allocator() function